use std::any::Any;
use std::sync::Arc;
use num_traits::AsPrimitive;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_ensure;
use vortex_error::vortex_panic;
use vortex_mask::Mask;
use crate::ArrayRef;
use crate::Canonical;
use crate::ExecutionCtx;
use crate::IntoArray;
use crate::array::ArrayView;
use crate::arrays::List;
use crate::arrays::ListArray;
use crate::arrays::ListView;
use crate::arrays::ListViewArray;
use crate::arrays::PrimitiveArray;
use crate::arrays::list::ListArrayExt;
use crate::arrays::listview::ListViewArrayExt;
use crate::builders::ArrayBuilder;
use crate::builders::DEFAULT_BUILDER_CAPACITY;
use crate::builders::LazyBitBufferBuilder;
use crate::builders::PrimitiveBuilder;
use crate::builders::builder_with_capacity;
use crate::dtype::DType;
use crate::dtype::IntegerPType;
use crate::dtype::Nullability;
use crate::dtype::Nullability::NonNullable;
use crate::match_each_integer_ptype;
use crate::scalar::ListScalar;
use crate::scalar::Scalar;
pub struct ListBuilder<O: IntegerPType> {
dtype: DType,
elements_builder: Box<dyn ArrayBuilder>,
offsets_builder: PrimitiveBuilder<O>,
nulls: LazyBitBufferBuilder,
}
impl<O: IntegerPType> ListBuilder<O> {
pub fn new(value_dtype: Arc<DType>, nullability: Nullability) -> Self {
Self::with_capacity(
value_dtype,
nullability,
DEFAULT_BUILDER_CAPACITY * 2,
DEFAULT_BUILDER_CAPACITY,
)
}
pub fn with_capacity(
value_dtype: Arc<DType>,
nullability: Nullability,
elements_capacity: usize,
capacity: usize,
) -> Self {
let elements_builder = builder_with_capacity(value_dtype.as_ref(), elements_capacity);
let mut offsets_builder = PrimitiveBuilder::<O>::with_capacity(NonNullable, capacity + 1);
offsets_builder.append_zero();
Self {
elements_builder,
offsets_builder,
nulls: LazyBitBufferBuilder::new(capacity),
dtype: DType::List(value_dtype, nullability),
}
}
pub fn append_array_as_list(
&mut self,
array: &ArrayRef,
ctx: &mut ExecutionCtx,
) -> VortexResult<()> {
vortex_ensure!(
array.dtype() == self.element_dtype(),
"Array dtype {:?} does not match list element dtype {:?}",
array.dtype(),
self.element_dtype()
);
self.elements_builder.reserve_exact(array.len());
array.append_to_builder(self.elements_builder.as_mut(), ctx)?;
self.nulls.append_non_null();
self.offsets_builder.append_value(
O::from_usize(self.elements_builder.len())
.vortex_expect("Failed to convert from usize to O"),
);
Ok(())
}
pub fn append_value(&mut self, value: ListScalar) -> VortexResult<()> {
match value.elements() {
None => {
if self.dtype.nullability() == NonNullable {
vortex_bail!("Cannot append null value to non-nullable list");
}
self.append_null();
}
Some(elements) => {
for scalar in elements {
self.elements_builder.append_scalar(&scalar)?;
}
self.nulls.append_non_null();
self.offsets_builder.append_value(
O::from_usize(self.elements_builder.len())
.vortex_expect("Failed to convert from usize to O"),
);
}
}
Ok(())
}
pub fn finish_into_list(&mut self) -> ListArray {
assert_eq!(
self.offsets_builder.len(),
self.nulls.len() + 1,
"offsets length must be one more than nulls length."
);
ListArray::try_new(
self.elements_builder.finish(),
self.offsets_builder.finish(),
self.nulls.finish_with_nullability(self.dtype.nullability()),
)
.vortex_expect("Buffer, offsets, and validity must have same length.")
}
pub fn element_dtype(&self) -> &DType {
let DType::List(element_dtype, _) = &self.dtype else {
vortex_panic!("`ListBuilder` has an incorrect dtype: {}", self.dtype);
};
element_dtype
}
}
fn extend_from_listview<O, OffsetType, SizeType>(
builder: &mut ListBuilder<O>,
new_elements: &ArrayRef,
new_offsets: &[OffsetType],
new_sizes: &[SizeType],
ctx: &mut ExecutionCtx,
) -> VortexResult<()>
where
O: IntegerPType,
OffsetType: IntegerPType,
SizeType: IntegerPType,
{
let num_lists = new_offsets.len();
debug_assert_eq!(num_lists, new_sizes.len());
let total_elements: usize = new_sizes.iter().map(|size| size.as_()).sum();
builder.elements_builder.reserve_exact(total_elements);
let mut curr_offset = builder.elements_builder.len();
builder.offsets_builder.reserve_exact(num_lists);
let mut offsets_range = builder.offsets_builder.uninit_range(num_lists);
for i in 0..new_offsets.len() {
let offset: usize = new_offsets[i].as_();
let size: usize = new_sizes[i].as_();
if size > 0 {
let list_elements = new_elements
.slice(offset..offset + size)
.vortex_expect("list builder slice");
list_elements.append_to_builder(builder.elements_builder.as_mut(), ctx)?;
curr_offset += size;
}
let new_offset = O::from_usize(curr_offset).vortex_expect("Failed to convert offset");
offsets_range.set_value(i, new_offset);
}
unsafe { offsets_range.finish() };
Ok(())
}
impl<O: IntegerPType> ArrayBuilder for ListBuilder<O> {
fn as_any(&self) -> &dyn Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
fn dtype(&self) -> &DType {
&self.dtype
}
fn len(&self) -> usize {
self.nulls.len()
}
fn append_zeros(&mut self, n: usize) {
let curr_len = self.elements_builder.len();
for _ in 0..n {
self.offsets_builder.append_value(
O::from_usize(curr_len).vortex_expect("Failed to convert from usize to <O>"),
)
}
self.nulls.append_n_non_nulls(n);
}
unsafe fn append_nulls_unchecked(&mut self, n: usize) {
let curr_len = self.elements_builder.len();
for _ in 0..n {
self.offsets_builder.append_value(
O::from_usize(curr_len).vortex_expect("Failed to convert from usize to <O>"),
)
}
self.nulls.append_n_nulls(n);
}
fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
vortex_ensure!(
scalar.dtype() == self.dtype(),
"ListBuilder expected scalar with dtype {}, got {}",
self.dtype(),
scalar.dtype()
);
self.append_value(scalar.as_list())
}
fn reserve_exact(&mut self, additional: usize) {
self.elements_builder.reserve_exact(additional);
self.offsets_builder.reserve_exact(additional);
self.nulls.reserve_exact(additional);
}
unsafe fn set_validity_unchecked(&mut self, validity: Mask) {
self.nulls = LazyBitBufferBuilder::from_validity_mask(validity);
}
fn finish(&mut self) -> ArrayRef {
self.finish_into_list().into_array()
}
fn finish_into_canonical(&mut self, ctx: &mut ExecutionCtx) -> Canonical {
let listview = self
.finish()
.execute::<ListViewArray>(ctx)
.vortex_expect("list builder should canonicalize to listview");
Canonical::List(listview)
}
fn append_list_array(
&mut self,
array: ArrayView<'_, List>,
ctx: &mut ExecutionCtx,
) -> VortexResult<()> {
if array.is_empty() {
return Ok(());
}
self.nulls
.append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?);
let num_lists = array.len();
let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
match_each_integer_ptype!(offsets.ptype(), |OffsetType| {
let offsets = offsets.as_slice::<OffsetType>();
let first: usize = offsets[0].as_();
let last: usize = offsets[num_lists].as_();
let elements_base = self.elements_builder.len();
if last > first {
self.elements_builder.reserve_exact(last - first);
array
.elements()
.slice(first..last)?
.append_to_builder(self.elements_builder.as_mut(), ctx)?;
}
self.offsets_builder.reserve_exact(num_lists);
let mut offsets_range = self.offsets_builder.uninit_range(num_lists);
for i in 0..num_lists {
let end: usize = offsets[i + 1].as_();
offsets_range.set_value(
i,
O::from_usize(end - first + elements_base)
.vortex_expect("Failed to convert offset"),
);
}
unsafe { offsets_range.finish() };
});
Ok(())
}
fn append_listview_array(
&mut self,
array: ArrayView<'_, ListView>,
ctx: &mut ExecutionCtx,
) -> VortexResult<()> {
if array.is_empty() {
return Ok(());
}
self.nulls
.append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?);
let elements = array.elements();
let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
let sizes = array.sizes().clone().execute::<PrimitiveArray>(ctx)?;
match_each_integer_ptype!(offsets.ptype(), |OffsetType| {
match_each_integer_ptype!(sizes.ptype(), |SizeType| {
extend_from_listview(
self,
elements,
offsets.as_slice::<OffsetType>(),
sizes.as_slice::<SizeType>(),
ctx,
)?
})
});
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use Nullability::NonNullable;
use Nullability::Nullable;
use vortex_buffer::buffer;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use crate::IntoArray;
use crate::array_session;
use crate::arrays::ChunkedArray;
use crate::arrays::ListViewArray;
use crate::arrays::PrimitiveArray;
use crate::arrays::list::ListArrayExt;
use crate::arrays::listview::ListViewArrayExt;
use crate::assert_arrays_eq;
use crate::builders::ArrayBuilder;
use crate::builders::ListViewBuilder;
use crate::builders::builder_with_capacity;
use crate::builders::list::ListArray;
use crate::builders::list::ListBuilder;
use crate::dtype::DType;
use crate::dtype::IntegerPType;
use crate::dtype::Nullability;
use crate::dtype::PType::I32;
use crate::executor::VortexSessionExecute;
use crate::scalar::Scalar;
use crate::validity::Validity;
#[test]
fn test_empty() {
let mut builder =
ListBuilder::<u32>::with_capacity(Arc::new(I32.into()), NonNullable, 0, 0);
let list = builder.finish();
assert_eq!(list.len(), 0);
}
#[test]
fn test_values() {
let dtype: Arc<DType> = Arc::new(I32.into());
let mut builder = ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0);
builder
.append_value(
Scalar::list(
Arc::clone(&dtype),
vec![1i32.into(), 2i32.into(), 3i32.into()],
NonNullable,
)
.as_list(),
)
.unwrap();
builder
.append_value(
Scalar::list(
dtype,
vec![4i32.into(), 5i32.into(), 6i32.into()],
NonNullable,
)
.as_list(),
)
.unwrap();
let list = builder.finish();
assert_eq!(list.len(), 2);
let mut ctx = array_session().create_execution_ctx();
let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3);
assert_eq!(list_array.list_elements_at(1).unwrap().len(), 3);
}
#[test]
fn test_append_empty_list() {
let dtype: Arc<DType> = Arc::new(I32.into());
let mut builder = ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0);
assert!(
builder
.append_value(Scalar::list_empty(dtype, NonNullable).as_list())
.is_ok()
)
}
#[test]
fn test_nullable_values() {
let dtype: Arc<DType> = Arc::new(I32.into());
let mut builder = ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), Nullable, 0, 0);
builder
.append_value(
Scalar::list(
Arc::clone(&dtype),
vec![1i32.into(), 2i32.into(), 3i32.into()],
NonNullable,
)
.as_list(),
)
.unwrap();
builder
.append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
.unwrap();
builder
.append_value(
Scalar::list(
dtype,
vec![4i32.into(), 5i32.into(), 6i32.into()],
NonNullable,
)
.as_list(),
)
.unwrap();
let list = builder.finish();
assert_eq!(list.len(), 3);
let mut ctx = array_session().create_execution_ctx();
let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3);
assert_eq!(list_array.list_elements_at(1).unwrap().len(), 0);
assert_eq!(list_array.list_elements_at(2).unwrap().len(), 3);
}
fn test_extend_builder_gen<O: IntegerPType>() {
let list = ListArray::from_iter_opt_slow::<O, _, _>(
[Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
Arc::new(I32.into()),
)
.unwrap()
.into_array();
assert_eq!(list.len(), 3);
let mut ctx = array_session().create_execution_ctx();
let mut builder = ListBuilder::<O>::with_capacity(Arc::new(I32.into()), Nullable, 18, 9);
list.append_to_builder(&mut builder, &mut ctx).unwrap();
list.append_to_builder(&mut builder, &mut ctx).unwrap();
list.slice(0..0)
.unwrap()
.append_to_builder(&mut builder, &mut ctx)
.unwrap();
list.slice(1..3)
.unwrap()
.append_to_builder(&mut builder, &mut ctx)
.unwrap();
let expected = ListArray::from_iter_opt_slow::<O, _, _>(
[
Some(vec![0, 1, 2]),
None,
Some(vec![4, 5]),
Some(vec![0, 1, 2]),
None,
Some(vec![4, 5]),
None,
Some(vec![4, 5]),
],
Arc::new(DType::Primitive(I32, NonNullable)),
)
.unwrap()
.into_array()
.execute::<ListViewArray>(&mut ctx)
.unwrap();
let actual = builder.finish_into_canonical(&mut ctx).into_listview();
assert_arrays_eq!(actual.elements(), expected.elements(), &mut ctx);
assert_arrays_eq!(actual.offsets(), expected.offsets(), &mut ctx);
assert!(
actual
.validity()
.vortex_expect("list validity should be derivable")
.mask_eq(
&expected
.validity()
.vortex_expect("list validity should be derivable"),
actual.len(),
&mut ctx,
)
.unwrap(),
);
}
#[test]
fn test_append_to_builder_any_list_builder() -> VortexResult<()> {
let mut ctx = array_session().create_execution_ctx();
let list = ListArray::from_iter_opt_slow::<u64, _, _>(
[Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
Arc::new(I32.into()),
)?
.into_array();
let listview = list
.clone()
.execute::<ListViewArray>(&mut ctx)?
.into_array();
let elem_dtype = || Arc::new(I32.into());
let mut listview_builder = builder_with_capacity(list.dtype(), list.len());
list.append_to_builder(listview_builder.as_mut(), &mut ctx)?;
assert_arrays_eq!(listview_builder.finish(), list, &mut ctx);
let mut lv_u32_u8 = ListViewBuilder::<u32, u8>::with_capacity(elem_dtype(), Nullable, 8, 4);
list.append_to_builder(&mut lv_u32_u8, &mut ctx)?;
assert_arrays_eq!(lv_u32_u8.finish(), list, &mut ctx);
let mut lv_i32_i16 =
ListViewBuilder::<i32, i16>::with_capacity(elem_dtype(), Nullable, 8, 4);
list.append_to_builder(&mut lv_i32_i16, &mut ctx)?;
assert_arrays_eq!(lv_i32_i16.finish(), list, &mut ctx);
let mut lv_u16_u16 =
ListViewBuilder::<u16, u16>::with_capacity(elem_dtype(), Nullable, 8, 4);
listview.append_to_builder(&mut lv_u16_u16, &mut ctx)?;
assert_arrays_eq!(lv_u16_u16.finish(), list, &mut ctx);
let mut list_builder = ListBuilder::<u32>::with_capacity(elem_dtype(), Nullable, 8, 4);
list.append_to_builder(&mut list_builder, &mut ctx)?;
assert_arrays_eq!(list_builder.finish(), list, &mut ctx);
let mut list_builder_i16 = ListBuilder::<i16>::with_capacity(elem_dtype(), Nullable, 8, 4);
listview.append_to_builder(&mut list_builder_i16, &mut ctx)?;
assert_arrays_eq!(list_builder_i16.finish(), list, &mut ctx);
Ok(())
}
#[test]
fn test_append_list_arrays_grow_builder() -> VortexResult<()> {
let mut ctx = array_session().create_execution_ctx();
let dtype: Arc<DType> = Arc::new(I32.into());
let lists: Vec<Option<Vec<i32>>> =
(0..100).map(|i| (i % 10 != 0).then(|| vec![i])).collect();
let source = ListArray::from_iter_opt_slow::<u32, _, _>(lists.clone(), Arc::clone(&dtype))?;
let expected = ListArray::from_iter_opt_slow::<u32, _, _>(
lists.iter().cloned().chain(lists.iter().cloned()),
Arc::clone(&dtype),
)?;
let mut builder = ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), Nullable, 0, 0);
builder.append_list_array(source.as_view(), &mut ctx)?;
builder.append_list_array(source.as_view(), &mut ctx)?;
assert_arrays_eq!(builder.finish(), expected, &mut ctx);
let source_listview = source.into_array().execute::<ListViewArray>(&mut ctx)?;
let mut builder = ListBuilder::<u32>::with_capacity(dtype, Nullable, 0, 0);
builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
assert_arrays_eq!(builder.finish(), expected, &mut ctx);
Ok(())
}
#[test]
fn test_extend_builder() {
test_extend_builder_gen::<i8>();
test_extend_builder_gen::<i16>();
test_extend_builder_gen::<i32>();
test_extend_builder_gen::<i64>();
test_extend_builder_gen::<u8>();
test_extend_builder_gen::<u16>();
test_extend_builder_gen::<u32>();
test_extend_builder_gen::<u64>();
}
#[test]
pub fn test_array_with_gap() {
let one_trailing_unused_element = ListArray::try_new(
buffer![1, 2, 3, 4].into_array(),
buffer![0, 3].into_array(),
Validity::NonNullable,
)
.unwrap();
let second_array = ListArray::try_new(
buffer![5, 6].into_array(),
buffer![0, 2].into_array(),
Validity::NonNullable,
)
.unwrap();
let chunked_list = ChunkedArray::try_new(
vec![
one_trailing_unused_element.clone().into_array(),
second_array.clone().into_array(),
],
DType::List(Arc::new(DType::Primitive(I32, NonNullable)), NonNullable),
);
let mut ctx = array_session().create_execution_ctx();
let canon_values = chunked_list
.unwrap()
.as_array()
.clone()
.execute::<ListViewArray>(&mut ctx)
.unwrap();
assert_eq!(
one_trailing_unused_element
.execute_scalar(0, &mut array_session().create_execution_ctx())
.unwrap(),
canon_values
.execute_scalar(0, &mut array_session().create_execution_ctx())
.unwrap()
);
assert_eq!(
second_array
.execute_scalar(0, &mut array_session().create_execution_ctx())
.unwrap(),
canon_values
.execute_scalar(1, &mut array_session().create_execution_ctx())
.unwrap()
);
}
#[test]
fn test_append_scalar() {
let dtype: Arc<DType> = Arc::new(I32.into());
let mut builder = ListBuilder::<u64>::with_capacity(Arc::clone(&dtype), Nullable, 20, 10);
let list_scalar1 =
Scalar::list(Arc::clone(&dtype), vec![1i32.into(), 2i32.into()], Nullable);
builder.append_scalar(&list_scalar1).unwrap();
let list_scalar2 = Scalar::list(
Arc::clone(&dtype),
vec![3i32.into(), 4i32.into(), 5i32.into()],
Nullable,
);
builder.append_scalar(&list_scalar2).unwrap();
let null_scalar = Scalar::null(DType::List(Arc::clone(&dtype), Nullable));
builder.append_scalar(&null_scalar).unwrap();
let array = builder.finish_into_list();
assert_eq!(array.len(), 3);
let mut ctx = array_session().create_execution_ctx();
let scalar0 = array.execute_scalar(0, &mut ctx).unwrap();
let list0 = scalar0.as_list();
assert_eq!(list0.len(), 2);
if let Some(list0_items) = list0.elements() {
assert_eq!(list0_items[0].as_primitive().typed_value::<i32>(), Some(1));
assert_eq!(list0_items[1].as_primitive().typed_value::<i32>(), Some(2));
}
let scalar1 = array.execute_scalar(1, &mut ctx).unwrap();
let list1 = scalar1.as_list();
assert_eq!(list1.len(), 3);
if let Some(list1_items) = list1.elements() {
assert_eq!(list1_items[0].as_primitive().typed_value::<i32>(), Some(3));
assert_eq!(list1_items[1].as_primitive().typed_value::<i32>(), Some(4));
assert_eq!(list1_items[2].as_primitive().typed_value::<i32>(), Some(5));
}
let scalar2 = array.execute_scalar(2, &mut ctx).unwrap();
let list2 = scalar2.as_list();
assert!(list2.is_null());
assert!(
array
.validity()
.vortex_expect("list validity should be derivable")
.execute_is_valid(0, &mut ctx)
.unwrap()
);
assert!(
array
.validity()
.vortex_expect("list validity should be derivable")
.execute_is_valid(1, &mut ctx)
.unwrap()
);
assert!(
!array
.validity()
.vortex_expect("list validity should be derivable")
.execute_is_valid(2, &mut ctx)
.unwrap()
);
let mut builder = ListBuilder::<u64>::with_capacity(dtype, NonNullable, 20, 10);
let wrong_scalar = Scalar::from(42i32);
assert!(builder.append_scalar(&wrong_scalar).is_err());
}
#[test]
fn test_append_array_as_list() {
let dtype: Arc<DType> = Arc::new(I32.into());
let mut ctx = array_session().create_execution_ctx();
let mut builder =
ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), NonNullable, 20, 10);
let arr1 = buffer![1i32, 2, 3].into_array();
builder.append_array_as_list(&arr1, &mut ctx).unwrap();
builder
.append_value(
Scalar::list(
Arc::clone(&dtype),
vec![10i32.into(), 11i32.into()],
NonNullable,
)
.as_list(),
)
.unwrap();
let arr2 = buffer![4i32, 5].into_array();
builder.append_array_as_list(&arr2, &mut ctx).unwrap();
let arr3 = buffer![0i32; 0].into_array();
builder.append_array_as_list(&arr3, &mut ctx).unwrap();
builder
.append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
.unwrap();
let list = builder.finish_into_list();
assert_eq!(list.len(), 5);
assert_arrays_eq!(
list.elements(),
PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]),
&mut ctx
);
assert_arrays_eq!(
list.offsets(),
PrimitiveArray::from_iter([0u32, 3, 5, 7, 7, 7]),
&mut ctx
);
let mut builder = ListBuilder::<u32>::with_capacity(dtype, NonNullable, 20, 10);
let wrong_dtype_arr = buffer![1i64, 2, 3].into_array();
assert!(
builder
.append_array_as_list(&wrong_dtype_arr, &mut ctx)
.is_err()
);
}
}