add list arb

spiraldb · Dec 18, 2024 · 9642d3a · 9642d3a
1 parent 0568e99
commit 9642d3a
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 13 deletions.
diff --git a/vortex-array/src/array/arbitrary.rs b/vortex-array/src/array/arbitrary.rs
@@ -5,7 +5,8 @@ use arrow_buffer::BooleanBuffer;
 use builders::ListBuilder;
 use vortex_dtype::{DType, NativePType, Nullability, PType};
 use vortex_error::{VortexExpect, VortexUnwrap};
-use vortex_scalar::ListScalar;
+use vortex_scalar::arbitrary::random_scalar;
+use vortex_scalar::{ListScalar, Scalar};
 
 use super::{BoolArray, ChunkedArray, ListArray, NullArray, PrimitiveArray, StructArray};
 use crate::array::{VarBinArray, VarBinViewArray};
@@ -86,11 +87,17 @@ fn random_array(u: &mut Unstructured, dtype: &DType, len: Option<usize>) -> Resu
                 }
                 DType::List(ldt, n) => {
                     let list_len = u.int_in_range(0..=20)?;
+                    let mut builder = ListBuilder::with_capacity(ldt.clone(), *n, 1);
                     for _ in 0..list_len {
-                        let value = random_array(u, ldt, None)?;
-                        let mut builder = ListBuilder::with_capacity(ldt.clone(), *n, 1);
-                        ListScalar::builder.append_elements(value)?;
-                        builder.finish()?;
+                        if u.arbitrary::<bool>() {
+                            let elem_len = u.int_in_range(0..=20)?;
+                            let elem = (0..elem_len)
+                                .map(|_| random_scalar(u, ldt))
+                                .collect::<Result<Vec<_>>>()?;
+                            builder.append_value(Scalar::list(ldt.clone(), elem).as_list())?;
+                        } else {
+                            builder.append_null();
+                        }
                     }
                     Ok(builder.finish()?)
                 }

diff --git a/vortex-array/src/array/list/mod.rs b/vortex-array/src/array/list/mod.rs
@@ -88,8 +88,6 @@ impl ListArray {
         )
     }
 
-    pub fn from_vec()
-
     pub fn validity(&self) -> Validity {
         self.metadata().validity.to_validity(|| {
             self.as_ref()

diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs
@@ -63,12 +63,6 @@ impl ListBuilder {
         }
     }
 
-    pub fn append_elements(&mut self, elements: ArrayData) -> VortexResult<()> {
-        elements.
-        self.value_builder.app
-        self.append_index(self.value_builder.len() as u64)
-    }
-
     fn append_index(&mut self, index: u64) -> VortexResult<()> {
         self.index_builder
             .append_scalar(&Scalar::from(index).cast(self.index_builder.dtype())?)

diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs
@@ -12,6 +12,7 @@ use std::any::Any;
 pub use binary::*;
 pub use bool::*;
 pub use extension::*;
+pub use list::*;
 pub use null::*;
 pub use primitive::*;
 pub use utf8::*;