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

vortex_array/arrays/
arbitrary.rs
