wip

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

encodings/runend/src/arbitrary.rs

@@ -9,6 +9,7 @@ use vortex_array::LEGACY_SESSION;
 use vortex_array::VortexSessionExecute;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::arbitrary::ArbitraryArray;
+use vortex_array::arrays::arbitrary::ArbitraryWith;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
@@ -46,14 +47,28 @@ impl ArbitraryRunEndArray {
         if num_runs == 0 {
             // Empty RunEndArray
             let ends = PrimitiveArray::from_iter(Vec::<u64>::new()).into_array();
-            let values = ArbitraryArray::arbitrary_with(u, Some(0), dtype)?.0;
+            let values = ArbitraryArray::arbitrary_with_config(
+                u,
+                &vortex_array::arrays::arbitrary::ArbitraryArrayConfig {
+                    dtype: Some(dtype.clone()),
+                    len: 0..=0,
+                },
+            )?
+            .0;
             let runend_array = RunEnd::try_new(ends, values, &mut ctx)
                 .vortex_expect("Empty RunEndArray creation should succeed");
             return Ok(ArbitraryRunEndArray(runend_array));
         }
 
         // Generate arbitrary values for each run
-        let values = ArbitraryArray::arbitrary_with(u, Some(num_runs), dtype)?.0;
+        let values = ArbitraryArray::arbitrary_with_config(
+            u,
+            &vortex_array::arrays::arbitrary::ArbitraryArrayConfig {
+                dtype: Some(dtype.clone()),
+                len: num_runs..=num_runs,
+            },
+        )?
+        .0;
 
         // Generate strictly increasing ends
         // Each end must be > previous end, and first end must be >= 1

vortex-array/src/arrays/arbitrary.rs

@@ -60,8 +60,13 @@ pub struct ArbitraryArrayConfig {
 
 impl<'a> Arbitrary<'a> for ArbitraryArray {
     fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
-        let dtype = u.arbitrary()?;
-        Self::arbitrary_with(u, None, &dtype)
+        Self::arbitrary_with_config(
+            u,
+            &ArbitraryArrayConfig {
+                dtype: None,
+                len: DEFAULT_ARRAY_LEN_RANGE,
+            },
+        )
     }
 }
 
@@ -84,12 +89,6 @@ impl<'a> ArbitraryWith<'a, ArbitraryArrayConfig> for ArbitraryArray {
     }
 }
 
-impl ArbitraryArray {
-    pub fn arbitrary_with(u: &mut Unstructured, len: Option<usize>, dtype: &DType) -> Result<Self> {
-        random_array(u, dtype, len).map(ArbitraryArray)
-    }
-}
-
 fn split_number_into_parts(n: usize, parts: usize) -> Vec<usize> {
     let reminder = n % parts;
     let division = (n - reminder) / parts;

vortex-array/src/arrays/dict/arbitrary.rs

@@ -13,6 +13,7 @@ use crate::ArrayRef;
 use crate::IntoArray;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::arbitrary::ArbitraryArray;
+use crate::arrays::arbitrary::ArbitraryWith;
 use crate::arrays::arbitrary::random_validity;
 use crate::dtype::DType;
 use crate::dtype::NativePType;
@@ -36,7 +37,14 @@ impl ArbitraryDictArray {
         // Generate the number of unique values (dictionary size)
         let values_len = u.int_in_range(1..=20)?;
         // Generate values array with the given dtype
-        let values = ArbitraryArray::arbitrary_with(u, Some(values_len), dtype)?.0;
+        let values = ArbitraryArray::arbitrary_with_config(
+            u,
+            &crate::arrays::arbitrary::ArbitraryArrayConfig {
+                dtype: Some(dtype.clone()),
+                len: values_len..=values_len,
+            },
+        )?
+        .0;
 
         // Generate codes that index into the values
         let codes_len = len.unwrap_or(u.int_in_range(0..=100)?);