use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use crate::ArrayRef;
use crate::Canonical;
use crate::ExecutionCtx;
use crate::IntoArray;
use crate::ToCanonical;
use crate::arrays::ExtensionArray;
use crate::arrays::FixedSizeListArray;
use crate::arrays::ListArray;
use crate::arrays::ListViewArray;
use crate::arrays::PrimitiveArray;
use crate::arrays::StructArray;
use crate::arrays::extension::ExtensionArrayExt;
use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
use crate::arrays::list::ListArrayExt;
use crate::arrays::listview::ListViewArrayExt;
use crate::arrays::listview::ListViewRebuildMode;
use crate::arrays::struct_::StructArrayExt;
use crate::builders::PrimitiveBuilder;
use crate::dtype::IntegerPType;
use crate::dtype::Nullability;
use crate::match_each_integer_ptype;
pub fn list_view_from_list(list: ListArray, ctx: &mut ExecutionCtx) -> VortexResult<ListViewArray> {
if list.is_empty() {
return Ok(Canonical::empty(list.dtype()).into_listview());
}
let list = list.reset_offsets(false).vortex_expect("This can't fail");
let list_offsets = list.offsets().clone();
let sizes = match_each_integer_ptype!(list_offsets.dtype().as_ptype(), |O| {
build_sizes_from_offsets::<O>(&list, ctx)?
});
debug_assert_eq!(list_offsets.len(), list.len() + 1);
let adjusted_offsets = list_offsets.slice(0..list.len())?;
Ok(unsafe {
ListViewArray::new_unchecked(
list.elements().clone(),
adjusted_offsets,
sizes,
list.validity()?,
)
.with_zero_copy_to_list(true)
})
}
fn build_sizes_from_offsets<O: IntegerPType>(
list: &ListArray,
ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
let len = list.len();
let mut sizes_builder = PrimitiveBuilder::<O>::with_capacity(Nullability::NonNullable, len);
let mut sizes_range = sizes_builder.uninit_range(len);
let offsets = list.offsets().clone().execute::<PrimitiveArray>(ctx)?;
let offsets_slice = offsets.as_slice::<O>();
debug_assert_eq!(len + 1, offsets_slice.len());
debug_assert!(offsets_slice.is_sorted());
for i in 0..len {
let size = offsets_slice[i + 1] - offsets_slice[i];
sizes_range.set_value(i, size);
}
unsafe {
sizes_range.finish();
}
Ok(sizes_builder.finish_into_primitive().into_array())
}
pub fn list_from_list_view(list_view: ListViewArray) -> VortexResult<ListArray> {
let zctl_array = list_view.rebuild(ListViewRebuildMode::MakeExact)?;
debug_assert!(zctl_array.is_zero_copy_to_list());
let list_offsets = match_each_integer_ptype!(zctl_array.offsets().dtype().as_ptype(), |O| {
unsafe { build_list_offsets_from_list_view::<O>(&zctl_array) }
});
Ok(unsafe {
ListArray::new_unchecked(
zctl_array.elements().clone(),
list_offsets,
zctl_array.validity()?,
)
})
}
unsafe fn build_list_offsets_from_list_view<O: IntegerPType>(
list_view: &ListViewArray,
) -> ArrayRef {
let len = list_view.len();
let mut offsets_builder =
PrimitiveBuilder::<O>::with_capacity(Nullability::NonNullable, len + 1);
let mut offsets_range = offsets_builder.uninit_range(len + 1);
let offsets = list_view.offsets().to_primitive();
let offsets_slice = offsets.as_slice::<O>();
debug_assert!(offsets_slice.is_sorted());
offsets_range.copy_from_slice(0, offsets_slice);
let final_offset = if len != 0 {
let last_offset = offsets_slice[len - 1];
let last_size = list_view.size_at(len - 1);
let last_size =
O::from_usize(last_size).vortex_expect("size somehow did not fit into offsets");
last_offset + last_size
} else {
O::zero()
};
offsets_range.set_value(len, final_offset);
unsafe {
offsets_range.finish();
}
offsets_builder.finish_into_primitive().into_array()
}
pub fn recursive_list_from_list_view(array: ArrayRef) -> VortexResult<ArrayRef> {
if !array.dtype().is_nested() {
return Ok(array);
}
let canonical = array.to_canonical()?;
Ok(match canonical {
Canonical::List(listview) => {
let converted_elements = recursive_list_from_list_view(listview.elements().clone())?;
debug_assert_eq!(converted_elements.len(), listview.elements().len());
let listview_with_converted_elements =
if !ArrayRef::ptr_eq(&converted_elements, listview.elements()) {
unsafe {
ListViewArray::new_unchecked(
converted_elements,
listview.offsets().clone(),
listview.sizes().clone(),
listview.validity()?,
)
.with_zero_copy_to_list(listview.is_zero_copy_to_list())
}
} else {
listview
};
let list_array = list_from_list_view(listview_with_converted_elements)?;
list_array.into_array()
}
Canonical::FixedSizeList(fixed_size_list) => {
let converted_elements =
recursive_list_from_list_view(fixed_size_list.elements().clone())?;
if !ArrayRef::ptr_eq(&converted_elements, fixed_size_list.elements()) {
FixedSizeListArray::try_new(
converted_elements,
fixed_size_list.list_size(),
fixed_size_list.validity()?,
fixed_size_list.len(),
)
.vortex_expect(
"FixedSizeListArray reconstruction should not fail with valid components",
)
.into_array()
} else {
fixed_size_list.into_array()
}
}
Canonical::Struct(struct_array) => {
let fields = struct_array.unmasked_fields();
let mut converted_fields = Vec::with_capacity(fields.len());
let mut any_changed = false;
for field in fields.iter() {
let converted_field = recursive_list_from_list_view(field.clone())?;
any_changed |= !ArrayRef::ptr_eq(&converted_field, field);
converted_fields.push(converted_field);
}
if any_changed {
StructArray::try_new(
struct_array.names().clone(),
converted_fields,
struct_array.len(),
struct_array.validity()?,
)
.vortex_expect("StructArray reconstruction should not fail with valid components")
.into_array()
} else {
struct_array.into_array()
}
}
Canonical::Extension(ext_array) => {
let converted_storage =
recursive_list_from_list_view(ext_array.storage_array().clone())?;
if !ArrayRef::ptr_eq(&converted_storage, ext_array.storage_array()) {
ExtensionArray::new(ext_array.ext_dtype().clone(), converted_storage).into_array()
} else {
ext_array.into_array()
}
}
_ => unreachable!(),
})
}
#[cfg(test)]
mod tests {
use vortex_buffer::buffer;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use super::super::tests::common::create_basic_listview;
use super::super::tests::common::create_empty_lists_listview;
use super::super::tests::common::create_nullable_listview;
use super::super::tests::common::create_overlapping_listview;
use super::recursive_list_from_list_view;
use crate::ArrayEq;
use crate::ArrayRef;
use crate::IntoArray;
use crate::LEGACY_SESSION;
use crate::Precision;
use crate::VortexSessionExecute;
use crate::arrays::BoolArray;
use crate::arrays::FixedSizeListArray;
use crate::arrays::ListArray;
use crate::arrays::ListViewArray;
use crate::arrays::PrimitiveArray;
use crate::arrays::StructArray;
use crate::arrays::VarBinViewArray;
use crate::arrays::list::ListArrayExt;
use crate::arrays::listview::ListViewArrayExt;
use crate::arrays::listview::list_from_list_view;
use crate::arrays::listview::list_view_from_list;
use crate::assert_arrays_eq;
use crate::dtype::FieldNames;
use crate::validity::Validity;
#[test]
fn test_list_to_listview_basic() -> VortexResult<()> {
let elements = buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array();
let offsets = buffer![0u32, 3, 5, 7, 10].into_array();
let list_array = ListArray::try_new(elements.clone(), offsets, Validity::NonNullable)?;
let mut ctx = LEGACY_SESSION.create_execution_ctx();
let list_view = list_view_from_list(list_array.clone(), &mut ctx)?;
assert_eq!(list_view.len(), 4);
assert_arrays_eq!(elements, list_view.elements().clone());
let expected_offsets = buffer![0u32, 3, 5, 7].into_array();
assert_arrays_eq!(expected_offsets, list_view.offsets().clone());
let expected_sizes = buffer![3u32, 2, 2, 3].into_array();
assert_arrays_eq!(expected_sizes, list_view.sizes().clone());
assert_arrays_eq!(list_array, list_view);
Ok(())
}
#[test]
fn test_listview_to_list_zero_copy() -> VortexResult<()> {
let list_view = create_basic_listview();
let list_array = list_from_list_view(list_view.clone())?;
assert_arrays_eq!(list_view.elements().clone(), list_array.elements().clone());
let list_array_offsets_without_last = list_array.offsets().slice(0..list_view.len())?;
assert_arrays_eq!(list_view.offsets().clone(), list_array_offsets_without_last);
assert_arrays_eq!(list_view, list_array);
Ok(())
}
#[test]
fn test_empty_array_conversions() -> VortexResult<()> {
let empty_elements = PrimitiveArray::from_iter::<[i32; 0]>([]).into_array();
let empty_offsets = buffer![0u32].into_array();
let empty_list = ListArray::try_new(empty_elements, empty_offsets, Validity::NonNullable)?;
let mut ctx = LEGACY_SESSION.create_execution_ctx();
let empty_list_view = list_view_from_list(empty_list.clone(), &mut ctx)?;
assert_eq!(empty_list_view.len(), 0);
let converted_back = list_from_list_view(empty_list_view)?;
assert_eq!(converted_back.len(), 0);
assert_eq!(empty_list.len(), converted_back.len());
Ok(())
}
#[test]
fn test_nullable_conversions() -> VortexResult<()> {
let elements = buffer![10i32, 20, 30, 40, 50].into_array();
let offsets = buffer![0u32, 2, 4, 5].into_array();
let validity = Validity::Array(BoolArray::from_iter(vec![true, false, true]).into_array());
let nullable_list = ListArray::try_new(elements, offsets, validity.clone())?;
let mut ctx = LEGACY_SESSION.create_execution_ctx();
let nullable_list_view = list_view_from_list(nullable_list.clone(), &mut ctx)?;
assert!(
nullable_list_view
.validity()
.vortex_expect("listview validity should be derivable")
.array_eq(&validity, Precision::Ptr)
);
assert_eq!(nullable_list_view.len(), 3);
let converted_back = list_from_list_view(nullable_list_view)?;
assert_arrays_eq!(nullable_list, converted_back);
Ok(())
}
#[test]
fn test_non_zero_copy_listview_to_list() -> VortexResult<()> {
let list_view = create_overlapping_listview();
let list_array = list_from_list_view(list_view.clone())?;
for i in 0..list_array.len() {
let start = list_array.offset_at(i)?;
let end = list_array.offset_at(i + 1)?;
assert!(end >= start, "Offsets should be monotonic after conversion");
}
assert_arrays_eq!(list_view, list_array);
Ok(())
}
#[test]
fn test_empty_sublists() -> VortexResult<()> {
let empty_lists_view = create_empty_lists_listview();
let list_array = list_from_list_view(empty_lists_view.clone())?;
assert_eq!(list_array.len(), 4);
for i in 0..list_array.len() {
assert_eq!(list_array.list_elements_at(i)?.len(), 0);
}
let mut ctx = LEGACY_SESSION.create_execution_ctx();
let converted_back = list_view_from_list(list_array, &mut ctx)?;
assert_arrays_eq!(empty_lists_view, converted_back);
Ok(())
}
#[test]
fn test_different_offset_types() -> VortexResult<()> {
let elements = buffer![1i32, 2, 3, 4, 5].into_array();
let i32_offsets = buffer![0i32, 2, 5].into_array();
let list_i32 =
ListArray::try_new(elements.clone(), i32_offsets.clone(), Validity::NonNullable)?;
let mut ctx = LEGACY_SESSION.create_execution_ctx();
let list_view_i32 = list_view_from_list(list_i32.clone(), &mut ctx)?;
assert_eq!(list_view_i32.offsets().dtype(), i32_offsets.dtype());
assert_eq!(list_view_i32.sizes().dtype(), i32_offsets.dtype());
let i64_offsets = buffer![0i64, 2, 5].into_array();
let list_i64 = ListArray::try_new(elements, i64_offsets.clone(), Validity::NonNullable)?;
let list_view_i64 = list_view_from_list(list_i64.clone(), &mut ctx)?;
assert_eq!(list_view_i64.offsets().dtype(), i64_offsets.dtype());
assert_eq!(list_view_i64.sizes().dtype(), i64_offsets.dtype());
assert_arrays_eq!(list_i32, list_view_i32);
assert_arrays_eq!(list_i64, list_view_i64);
Ok(())
}
#[test]
fn test_round_trip_conversions() -> VortexResult<()> {
let mut ctx = LEGACY_SESSION.create_execution_ctx();
let original = create_basic_listview();
let to_list = list_from_list_view(original.clone())?;
let back_to_view = list_view_from_list(to_list, &mut ctx)?;
assert_arrays_eq!(original, back_to_view);
let nullable = create_nullable_listview();
let nullable_to_list = list_from_list_view(nullable.clone())?;
let nullable_back = list_view_from_list(nullable_to_list, &mut ctx)?;
assert_arrays_eq!(nullable, nullable_back);
let overlapping = create_overlapping_listview();
let overlapping_to_list = list_from_list_view(overlapping.clone())?;
let overlapping_back = list_view_from_list(overlapping_to_list, &mut ctx)?;
assert_arrays_eq!(overlapping, overlapping_back);
Ok(())
}
#[test]
fn test_single_element_lists() -> VortexResult<()> {
let elements = buffer![100i32, 200, 300].into_array();
let offsets = buffer![0u32, 1, 2, 3].into_array();
let single_elem_list = ListArray::try_new(elements, offsets, Validity::NonNullable)?;
let mut ctx = LEGACY_SESSION.create_execution_ctx();
let list_view = list_view_from_list(single_elem_list.clone(), &mut ctx)?;
assert_eq!(list_view.len(), 3);
let expected_sizes = buffer![1u32, 1, 1].into_array();
assert_arrays_eq!(expected_sizes, list_view.sizes().clone());
let converted_back = list_from_list_view(list_view)?;
assert_arrays_eq!(single_elem_list, converted_back);
Ok(())
}
#[test]
fn test_mixed_empty_and_non_empty_lists() -> VortexResult<()> {
let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array();
let offsets = buffer![0u32, 2, 2, 3, 3, 6].into_array();
let mixed_list = ListArray::try_new(elements, offsets, Validity::NonNullable)?;
let mut ctx = LEGACY_SESSION.create_execution_ctx();
let list_view = list_view_from_list(mixed_list.clone(), &mut ctx)?;
assert_eq!(list_view.len(), 5);
let expected_sizes = buffer![2u32, 0, 1, 0, 3].into_array();
assert_arrays_eq!(expected_sizes, list_view.sizes().clone());
let converted_back = list_from_list_view(list_view)?;
assert_arrays_eq!(mixed_list, converted_back);
Ok(())
}
#[test]
fn test_recursive_simple_listview() -> VortexResult<()> {
let list_view = create_basic_listview();
let result = recursive_list_from_list_view(list_view.clone().into_array())?;
assert_eq!(result.len(), list_view.len());
assert_arrays_eq!(list_view.into_array(), result);
Ok(())
}
#[test]
fn test_recursive_nested_listview() -> VortexResult<()> {
let inner_elements = buffer![1i32, 2, 3].into_array();
let inner_offsets = buffer![0u32, 2].into_array();
let inner_sizes = buffer![2u32, 1].into_array();
let inner_listview = unsafe {
ListViewArray::new_unchecked(
inner_elements,
inner_offsets,
inner_sizes,
Validity::NonNullable,
)
.with_zero_copy_to_list(true)
};
let outer_offsets = buffer![0u32, 1].into_array();
let outer_sizes = buffer![1u32, 1].into_array();
let outer_listview = unsafe {
ListViewArray::new_unchecked(
inner_listview.into_array(),
outer_offsets,
outer_sizes,
Validity::NonNullable,
)
.with_zero_copy_to_list(true)
};
let result = recursive_list_from_list_view(outer_listview.clone().into_array())?;
assert_eq!(result.len(), 2);
assert_arrays_eq!(outer_listview.into_array(), result);
Ok(())
}
#[test]
fn test_recursive_struct_with_listview_fields() -> VortexResult<()> {
let listview_field = create_basic_listview().into_array();
let primitive_field = buffer![10i32, 20, 30, 40].into_array();
let struct_array = StructArray::try_new(
FieldNames::from(["lists", "values"]),
vec![listview_field, primitive_field],
4,
Validity::NonNullable,
)?;
let result = recursive_list_from_list_view(struct_array.clone().into_array())?;
assert_eq!(result.len(), 4);
assert_arrays_eq!(struct_array.into_array(), result);
Ok(())
}
#[test]
fn test_recursive_fixed_size_list_with_listview_elements() -> VortexResult<()> {
let lv1_elements = buffer![1i32, 2].into_array();
let lv1_offsets = buffer![0u32].into_array();
let lv1_sizes = buffer![2u32].into_array();
let lv1 = unsafe {
ListViewArray::new_unchecked(
lv1_elements,
lv1_offsets,
lv1_sizes,
Validity::NonNullable,
)
.with_zero_copy_to_list(true)
};
let lv2_elements = buffer![3i32, 4].into_array();
let lv2_offsets = buffer![0u32].into_array();
let lv2_sizes = buffer![2u32].into_array();
let lv2 = unsafe {
ListViewArray::new_unchecked(
lv2_elements,
lv2_offsets,
lv2_sizes,
Validity::NonNullable,
)
.with_zero_copy_to_list(true)
};
let dtype = lv1.dtype().clone();
let chunked_listviews =
crate::arrays::ChunkedArray::try_new(vec![lv1.into_array(), lv2.into_array()], dtype)?;
let fixed_list =
FixedSizeListArray::new(chunked_listviews.into_array(), 1, Validity::NonNullable, 2);
let result = recursive_list_from_list_view(fixed_list.clone().into_array())?;
assert_eq!(result.len(), 2);
assert_arrays_eq!(fixed_list.into_array(), result);
Ok(())
}
#[test]
fn test_recursive_deep_nesting() -> VortexResult<()> {
let innermost_elements = buffer![1i32, 2, 3].into_array();
let innermost_offsets = buffer![0u32, 2].into_array();
let innermost_sizes = buffer![2u32, 1].into_array();
let innermost_listview = unsafe {
ListViewArray::new_unchecked(
innermost_elements,
innermost_offsets,
innermost_sizes,
Validity::NonNullable,
)
.with_zero_copy_to_list(true)
};
let struct_array = StructArray::try_new(
FieldNames::from(["inner_lists"]),
vec![innermost_listview.into_array()],
2,
Validity::NonNullable,
)?;
let outer_offsets = buffer![0u32, 1].into_array();
let outer_sizes = buffer![1u32, 1].into_array();
let outer_listview = unsafe {
ListViewArray::new_unchecked(
struct_array.into_array(),
outer_offsets,
outer_sizes,
Validity::NonNullable,
)
.with_zero_copy_to_list(true)
};
let result = recursive_list_from_list_view(outer_listview.clone().into_array())?;
assert_eq!(result.len(), 2);
assert_arrays_eq!(outer_listview.into_array(), result);
Ok(())
}
#[test]
fn test_recursive_primitive_unchanged() -> VortexResult<()> {
let prim = buffer![1i32, 2, 3].into_array();
let prim_clone = prim.clone();
let result = recursive_list_from_list_view(prim)?;
assert!(ArrayRef::ptr_eq(&result, &prim_clone));
Ok(())
}
#[test]
fn test_recursive_mixed_listview_and_list() -> VortexResult<()> {
let listview = create_basic_listview();
let list = list_from_list_view(listview.clone())?;
let struct_array = StructArray::try_new(
FieldNames::from(["listview_field", "list_field"]),
vec![listview.into_array(), list.into_array()],
4,
Validity::NonNullable,
)?;
let result = recursive_list_from_list_view(struct_array.clone().into_array())?;
assert_eq!(result.len(), 4);
assert_arrays_eq!(struct_array.into_array(), result);
Ok(())
}
#[test]
fn test_empty_listview_to_list_without_zctl_flag() -> VortexResult<()> {
let elements = VarBinViewArray::from_iter_str(Vec::<&str>::new()).into_array();
let offsets = PrimitiveArray::from_iter(Vec::<i16>::new()).into_array();
let sizes = PrimitiveArray::from_iter(Vec::<i16>::new()).into_array();
let list_view = ListViewArray::try_new(elements, offsets, sizes, Validity::AllValid)?;
assert!(!list_view.is_zero_copy_to_list());
let list_array = list_from_list_view(list_view)?;
assert_eq!(list_array.len(), 0);
Ok(())
}
}