vortex_array/arrays/
arbitrary.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::iter;
5use std::ops::RangeInclusive;
6use std::sync::Arc;
7
8use arbitrary::Arbitrary;
9use arbitrary::Error::IncorrectFormat;
10use arbitrary::Result;
11use arbitrary::Unstructured;
12use vortex_buffer::BitBuffer;
13use vortex_buffer::Buffer;
14use vortex_error::VortexExpect;
15
16use crate::ArrayRef;
17use crate::IntoArray;
18use crate::arrays::BoolArray;
19use crate::arrays::ChunkedArray;
20use crate::arrays::NullArray;
21use crate::arrays::Primitive;
22use crate::arrays::PrimitiveArray;
23use crate::arrays::StructArray;
24use crate::arrays::VarBinArray;
25use crate::arrays::VarBinViewArray;
26use crate::arrays::primitive::PrimitiveArrayExt;
27use crate::builders::ArrayBuilder;
28use crate::builders::DecimalBuilder;
29use crate::builders::FixedSizeListBuilder;
30use crate::builders::ListViewBuilder;
31use crate::builders::MapBuilder;
32use crate::dtype::DType;
33use crate::dtype::IntegerPType;
34use crate::dtype::MapDType;
35use crate::dtype::NativePType;
36use crate::dtype::Nullability;
37use crate::dtype::OffsetBuilderPType;
38use crate::dtype::PType;
39use crate::match_each_decimal_value_type;
40use crate::scalar::Scalar;
41use crate::scalar::arbitrary::random_scalar;
42use crate::validity::Validity;
43
44/// A wrapper type to implement `Arbitrary` for `ArrayRef`.
45#[derive(Clone, Debug)]
46pub struct ArbitraryArray(pub ArrayRef);
47
48/// Trait for generating arbitrary values with a caller-provided configuration.
49pub trait ArbitraryWith<'a, C>: Sized {
50    /// Generate an arbitrary value using the provided configuration.
51    fn arbitrary_with_config(u: &mut Unstructured<'a>, config: &C) -> Result<Self>;
52}
53
54/// Configuration for arbitrary array generation.
55#[derive(Clone, Debug)]
56pub struct ArbitraryArrayConfig {
57    /// Fixed dtype, or `None` to generate one from [`Unstructured`].
58    pub dtype: Option<DType>,
59    /// Inclusive range for the total array length.
60    pub len: RangeInclusive<usize>,
61}
62
63impl<'a> ArbitraryWith<'a, ArbitraryArrayConfig> for ArbitraryArray {
64    fn arbitrary_with_config(
65        u: &mut Unstructured<'a>,
66        config: &ArbitraryArrayConfig,
67    ) -> Result<Self> {
68        if config.len.is_empty() {
69            return Err(IncorrectFormat);
70        }
71
72        let dtype = match &config.dtype {
73            Some(dtype) => dtype.clone(),
74            None => u.arbitrary()?,
75        };
76        let len = u.int_in_range(config.len.clone())?;
77
78        random_array(u, &dtype, Some(len)).map(ArbitraryArray)
79    }
80}
81
82fn split_number_into_parts(n: usize, parts: usize) -> Vec<usize> {
83    let reminder = n % parts;
84    let division = (n - reminder) / parts;
85    iter::repeat_n(division, parts - reminder)
86        .chain(iter::repeat_n(division + 1, reminder))
87        .collect()
88}
89
90/// Creates a random array with a random number of chunks.
91fn random_array(u: &mut Unstructured, dtype: &DType, len: Option<usize>) -> Result<ArrayRef> {
92    let num_chunks = u.int_in_range(1..=3)?;
93    let chunk_lens = len.map(|l| split_number_into_parts(l, num_chunks));
94    let mut chunks = (0..num_chunks)
95        .map(|i| {
96            let chunk_len = chunk_lens.as_ref().map(|c| c[i]);
97            random_array_chunk(u, dtype, chunk_len)
98        })
99        .collect::<Result<Vec<_>>>()?;
100
101    if chunks.len() == 1 {
102        Ok(chunks.remove(0))
103    } else {
104        let dtype = chunks[0].dtype().clone();
105        Ok(ChunkedArray::try_new(chunks, dtype)
106            .vortex_expect("operation should succeed in arbitrary impl")
107            .into_array())
108    }
109}
110
111/// Creates a random array chunk.
112fn random_array_chunk(
113    u: &mut Unstructured<'_>,
114    dtype: &DType,
115    chunk_len: Option<usize>,
116) -> Result<ArrayRef> {
117    match dtype {
118        DType::Null => Ok(NullArray::new(
119            chunk_len
120                .map(Ok)
121                .unwrap_or_else(|| u.int_in_range(0..=100))?,
122        )
123        .into_array()),
124        DType::Bool(n) => random_bool(u, *n, chunk_len),
125        DType::Primitive(ptype, n) => match ptype {
126            PType::U8 => random_primitive::<u8>(u, *n, chunk_len),
127            PType::U16 => random_primitive::<u16>(u, *n, chunk_len),
128            PType::U32 => random_primitive::<u32>(u, *n, chunk_len),
129            PType::U64 => random_primitive::<u64>(u, *n, chunk_len),
130            PType::I8 => random_primitive::<i8>(u, *n, chunk_len),
131            PType::I16 => random_primitive::<i16>(u, *n, chunk_len),
132            PType::I32 => random_primitive::<i32>(u, *n, chunk_len),
133            PType::I64 => random_primitive::<i64>(u, *n, chunk_len),
134            PType::F16 => {
135                let prim = random_primitive::<u16>(u, *n, chunk_len)?
136                    .as_::<Primitive>()
137                    .reinterpret_cast(PType::F16)
138                    .into_array();
139                Ok(prim)
140            }
141            PType::F32 => random_primitive::<f32>(u, *n, chunk_len),
142            PType::F64 => random_primitive::<f64>(u, *n, chunk_len),
143        },
144        d @ DType::Decimal(decimal, n) => {
145            let elem_len = chunk_len.unwrap_or(u.int_in_range(0..=20)?);
146            match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(decimal), |D| {
147                let mut builder = DecimalBuilder::new::<D>(*decimal, *n);
148                for _i in 0..elem_len {
149                    let random_decimal = random_scalar(u, d)?;
150                    builder.append_scalar(&random_decimal).vortex_expect(
151                        "was somehow unable to append a decimal to a decimal builder",
152                    );
153                }
154                Ok(builder.finish())
155            })
156        }
157        DType::Utf8(n) => random_string(u, *n, chunk_len),
158        DType::Binary(n) => random_bytes(u, *n, chunk_len),
159        DType::List(elem_dtype, null) => random_list(u, elem_dtype, *null, chunk_len),
160        DType::FixedSizeList(elem_dtype, list_size, null) => {
161            random_fixed_size_list(u, elem_dtype, *list_size, *null, chunk_len)
162        }
163        DType::Map(map_dtype, nullability) => {
164            random_map(u, map_dtype.clone(), *nullability, chunk_len)
165        }
166        DType::Struct(sdt, n) => {
167            let first_array = sdt
168                .fields()
169                .next()
170                .map(|d| random_array(u, &d, chunk_len))
171                .transpose()?;
172            let resolved_len = first_array
173                .as_ref()
174                .map(|a| a.len())
175                .or(chunk_len)
176                .map(Ok)
177                .unwrap_or_else(|| u.int_in_range(0..=100))?;
178            let children = first_array
179                .into_iter()
180                .map(Ok)
181                .chain(
182                    sdt.fields()
183                        .skip(1)
184                        .map(|d| random_array(u, &d, Some(resolved_len))),
185                )
186                .collect::<Result<Vec<_>>>()?;
187            Ok(StructArray::try_new(
188                sdt.names().clone(),
189                children,
190                resolved_len,
191                random_validity(u, *n, resolved_len)?,
192            )
193            .vortex_expect("operation should succeed in arbitrary impl")
194            .into_array())
195        }
196        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
197        DType::Variant(_) => {
198            unimplemented!("Variant arrays are not implemented")
199        }
200        DType::Extension(..) => {
201            unimplemented!("Extension arrays are not implemented")
202        }
203    }
204}
205
206fn random_map(
207    u: &mut Unstructured,
208    map_dtype: MapDType,
209    nullability: Nullability,
210    chunk_len: Option<usize>,
211) -> Result<ArrayRef> {
212    let array_length = chunk_len.unwrap_or(u.int_in_range(0..=20)?);
213    let key_dtype = map_dtype.key_dtype();
214    let value_dtype = map_dtype.value_dtype();
215    let dtype = DType::Map(map_dtype.clone(), nullability);
216    let mut builder = MapBuilder::<u64, u64>::with_capacity(map_dtype, nullability, array_length);
217
218    for _ in 0..array_length {
219        if nullability == Nullability::Nullable && u.arbitrary::<bool>()? {
220            builder.append_null();
221        } else {
222            let entry_count = u.int_in_range(0..=20)?;
223            let entries = (0..entry_count)
224                .map(|_| {
225                    let key = random_scalar(u, &key_dtype)?;
226                    let value = random_scalar(u, &value_dtype)?;
227                    Ok((key, value))
228                })
229                .collect::<Result<Vec<_>>>()?;
230            let scalar = Scalar::try_map(dtype.clone(), entries)
231                .vortex_expect("generated map scalar should be valid");
232            builder
233                .append_scalar(&scalar)
234                .vortex_expect("generated map scalar should append");
235        }
236    }
237
238    Ok(builder.finish_into_map().into_array())
239}
240
241/// Creates a random fixed-size list array.
242///
243/// If the `chunk_len` is specified, the length of the array will be equal to the chunk length.
244fn random_fixed_size_list(
245    u: &mut Unstructured,
246    elem_dtype: &Arc<DType>,
247    list_size: u32,
248    null: Nullability,
249    chunk_len: Option<usize>,
250) -> Result<ArrayRef> {
251    let array_length = chunk_len.unwrap_or(u.int_in_range(0..=20)?);
252
253    let mut builder =
254        FixedSizeListBuilder::with_capacity(Arc::clone(elem_dtype), list_size, null, array_length);
255
256    for _ in 0..array_length {
257        if null == Nullability::Nullable && u.arbitrary::<bool>()? {
258            builder.append_null();
259        } else {
260            builder
261                .append_value(random_list_scalar(u, elem_dtype, list_size, null)?.as_list())
262                .vortex_expect("can append value");
263        }
264    }
265
266    Ok(builder.finish())
267}
268
269/// Creates a random list array.
270///
271/// If the `chunk_len` is specified, the length of the array will be equal to the chunk length.
272fn random_list(
273    u: &mut Unstructured,
274    elem_dtype: &Arc<DType>,
275    null: Nullability,
276    chunk_len: Option<usize>,
277) -> Result<ArrayRef> {
278    let array_length = chunk_len.unwrap_or(u.int_in_range(0..=20)?);
279    // Worst-case total elements: each list can have up to 20 elements.
280    let max_total_elements = array_length as u64 * 20;
281
282    match u.int_in_range(0..=3)? {
283        0 if i32::max_value_as_u64() >= max_total_elements => {
284            random_list_with_offset_type::<i32>(u, elem_dtype, null, array_length)
285        }
286        1 if u32::max_value_as_u64() >= max_total_elements => {
287            random_list_with_offset_type::<u32>(u, elem_dtype, null, array_length)
288        }
289        // i64 and u64 always fit; also the fallback for when narrower types don't.
290        _ => {
291            if u.arbitrary::<bool>()? {
292                random_list_with_offset_type::<i64>(u, elem_dtype, null, array_length)
293            } else {
294                random_list_with_offset_type::<u64>(u, elem_dtype, null, array_length)
295            }
296        }
297    }
298}
299
300/// Creates a random list array with the given [`OffsetBuilderPType`] for the internal offsets child.
301fn random_list_with_offset_type<O: OffsetBuilderPType>(
302    u: &mut Unstructured,
303    elem_dtype: &Arc<DType>,
304    null: Nullability,
305    array_length: usize,
306) -> Result<ArrayRef> {
307    let mut builder =
308        ListViewBuilder::<O, O>::with_capacity(Arc::clone(elem_dtype), null, array_length, 10);
309
310    for _ in 0..array_length {
311        if null == Nullability::Nullable && u.arbitrary::<bool>()? {
312            builder.append_null();
313        } else {
314            let list_size = u.int_in_range(0..=20)?;
315            builder
316                .append_value(random_list_scalar(u, elem_dtype, list_size, null)?.as_list())
317                .vortex_expect("can append value");
318        }
319    }
320
321    Ok(builder.finish())
322}
323
324/// Creates a random list scalar with the specified list size.
325fn random_list_scalar(
326    u: &mut Unstructured,
327    elem_dtype: &Arc<DType>,
328    list_size: u32,
329    null: Nullability,
330) -> Result<Scalar> {
331    let elems = (0..list_size)
332        .map(|_| random_scalar(u, elem_dtype))
333        .collect::<Result<Vec<_>>>()?;
334    Ok(Scalar::list(Arc::clone(elem_dtype), elems, null))
335}
336
337fn random_string(
338    u: &mut Unstructured,
339    nullability: Nullability,
340    len: Option<usize>,
341) -> Result<ArrayRef> {
342    match nullability {
343        Nullability::NonNullable => {
344            let v = arbitrary_vec_of_len::<String>(u, len)?;
345            Ok(match u.int_in_range(0..=1)? {
346                0 => VarBinArray::from_vec(v, DType::Utf8(Nullability::NonNullable)).into_array(),
347                1 => VarBinViewArray::from_iter_str(v).into_array(),
348                _ => unreachable!(),
349            })
350        }
351        Nullability::Nullable => {
352            let v = arbitrary_vec_of_len::<Option<String>>(u, len)?;
353            Ok(match u.int_in_range(0..=1)? {
354                0 => VarBinArray::from_iter(v, DType::Utf8(Nullability::Nullable)).into_array(),
355                1 => VarBinViewArray::from_iter_nullable_str(v).into_array(),
356                _ => unreachable!(),
357            })
358        }
359    }
360}
361
362fn random_bytes(
363    u: &mut Unstructured,
364    nullability: Nullability,
365    len: Option<usize>,
366) -> Result<ArrayRef> {
367    match nullability {
368        Nullability::NonNullable => {
369            let v = arbitrary_vec_of_len::<Vec<u8>>(u, len)?;
370            Ok(match u.int_in_range(0..=1)? {
371                0 => VarBinArray::from_vec(v, DType::Binary(Nullability::NonNullable)).into_array(),
372                1 => VarBinViewArray::from_iter_bin(v).into_array(),
373                _ => unreachable!(),
374            })
375        }
376        Nullability::Nullable => {
377            let v = arbitrary_vec_of_len::<Option<Vec<u8>>>(u, len)?;
378            Ok(match u.int_in_range(0..=1)? {
379                0 => VarBinArray::from_iter(v, DType::Binary(Nullability::Nullable)).into_array(),
380                1 => VarBinViewArray::from_iter_nullable_bin(v).into_array(),
381                _ => unreachable!(),
382            })
383        }
384    }
385}
386
387fn random_primitive<'a, T: Arbitrary<'a> + NativePType>(
388    u: &mut Unstructured<'a>,
389    nullability: Nullability,
390    len: Option<usize>,
391) -> Result<ArrayRef> {
392    let v = arbitrary_vec_of_len::<T>(u, len)?;
393    let validity = random_validity(u, nullability, v.len())?;
394    Ok(PrimitiveArray::new(Buffer::copy_from(v), validity).into_array())
395}
396
397fn random_bool(
398    u: &mut Unstructured,
399    nullability: Nullability,
400    len: Option<usize>,
401) -> Result<ArrayRef> {
402    let v = arbitrary_vec_of_len(u, len)?;
403    let validity = random_validity(u, nullability, v.len())?;
404    Ok(BoolArray::new(BitBuffer::from(v), validity).into_array())
405}
406
407pub fn random_validity(
408    u: &mut Unstructured,
409    nullability: Nullability,
410    len: usize,
411) -> Result<Validity> {
412    match nullability {
413        Nullability::NonNullable => Ok(Validity::NonNullable),
414        Nullability::Nullable => Ok(match u.int_in_range(0..=2)? {
415            0 => Validity::AllValid,
416            1 => Validity::AllInvalid,
417            2 => Validity::from_iter(arbitrary_vec_of_len::<bool>(u, Some(len))?),
418            _ => unreachable!(),
419        }),
420    }
421}
422
423fn arbitrary_vec_of_len<'a, T: Arbitrary<'a>>(
424    u: &mut Unstructured<'a>,
425    len: Option<usize>,
426) -> Result<Vec<T>> {
427    len.map(|l| (0..l).map(|_| T::arbitrary(u)).collect::<Result<Vec<_>>>())
428        .unwrap_or_else(|| Vec::<T>::arbitrary(u))
429}
vortex_array/arrays/arbitrary.rs

vortex_array/arrays/
arbitrary.rs
