vortex_array/scalar/
arbitrary.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Arbitrary scalar value generation.
5//!
6//! This module provides functions to generate arbitrary scalar values of various data types.
7//! It is used by the fuzzer to test the correctness of the scalar value implementation.
8
9use std::iter;
10
11use arbitrary::Result;
12use arbitrary::Unstructured;
13use vortex_buffer::BufferString;
14use vortex_buffer::ByteBuffer;
15use vortex_error::VortexExpect;
16
17use crate::dtype::DType;
18use crate::dtype::DecimalDType;
19use crate::dtype::NativeDecimalType;
20use crate::dtype::PType;
21use crate::dtype::half::f16;
22use crate::match_each_decimal_value_type;
23use crate::scalar::DecimalValue;
24use crate::scalar::PValue;
25use crate::scalar::Scalar;
26use crate::scalar::ScalarValue;
27
28/// Generates an arbitrary [`Scalar`] of the given [`DType`].
29///
30/// # Errors
31///
32/// Returns an error if the underlying arbitrary generation fails.
33pub fn random_scalar(u: &mut Unstructured, dtype: &DType) -> Result<Scalar> {
34    // For nullable types, return null ~25% of the time. This is just to make sure we don't generate
35    // too few nulls.
36    if dtype.is_nullable() && u.ratio(1, 4)? {
37        return Ok(Scalar::null(dtype.clone()));
38    }
39
40    Ok(match dtype {
41        DType::Null => Scalar::null(dtype.clone()),
42        DType::Bool(_) => Scalar::try_new(dtype.clone(), Some(ScalarValue::Bool(u.arbitrary()?)))
43            .vortex_expect("unable to construct random `Scalar`_"),
44        DType::Primitive(p, _) => Scalar::try_new(
45            dtype.clone(),
46            Some(ScalarValue::Primitive(random_pvalue(u, p)?)),
47        )
48        .vortex_expect("unable to construct random `Scalar`_"),
49        DType::Decimal(decimal_type, _) => {
50            Scalar::try_new(dtype.clone(), Some(random_decimal(u, decimal_type)?))
51                .vortex_expect("unable to construct random `Scalar`_")
52        }
53        DType::Utf8(_) => Scalar::try_new(
54            dtype.clone(),
55            Some(ScalarValue::Utf8(BufferString::from(
56                u.arbitrary::<String>()?,
57            ))),
58        )
59        .vortex_expect("unable to construct random `Scalar`_"),
60        DType::Binary(_) => Scalar::try_new(
61            dtype.clone(),
62            Some(ScalarValue::Binary(ByteBuffer::from(
63                u.arbitrary::<Vec<u8>>()?,
64            ))),
65        )
66        .vortex_expect("unable to construct random `Scalar`_"),
67        DType::List(edt, _) => Scalar::try_new(
68            dtype.clone(),
69            Some(ScalarValue::Tuple(
70                iter::from_fn(|| {
71                    // Generate elements with 1/4 probability.
72                    u.arbitrary()
73                        .unwrap_or(false)
74                        .then(|| random_scalar(u, edt).map(|s| s.into_value()))
75                })
76                .collect::<Result<Vec<_>>>()?,
77            )),
78        )
79        .vortex_expect("unable to construct random `Scalar`_"),
80        DType::FixedSizeList(edt, size, _) => Scalar::try_new(
81            dtype.clone(),
82            Some(ScalarValue::Tuple(
83                (0..*size)
84                    .map(|_| random_scalar(u, edt).map(|s| s.into_value()))
85                    .collect::<Result<Vec<_>>>()?,
86            )),
87        )
88        .vortex_expect("unable to construct random `Scalar`_"),
89        DType::Struct(sdt, _) => Scalar::try_new(
90            dtype.clone(),
91            Some(ScalarValue::Tuple(
92                sdt.fields()
93                    .map(|d| random_scalar(u, &d).map(|s| s.into_value()))
94                    .collect::<Result<Vec<_>>>()?,
95            )),
96        )
97        .vortex_expect("unable to construct random `Scalar`_"),
98        DType::Union(variants, nullability) => {
99            let child_index = u.choose_index(variants.len())?;
100
101            let child_dtype = variants
102                .variant_by_index(child_index)
103                .vortex_expect("chosen union child index must be valid");
104            let child = random_scalar(u, &child_dtype)?;
105
106            Scalar::union(
107                variants.clone(),
108                variants.child_index_to_tag(child_index),
109                child,
110                *nullability,
111            )
112            .vortex_expect("generated union scalar must be valid")
113        }
114        DType::Variant(_) => todo!(),
115        DType::Extension(..) => {
116            unreachable!("Can't yet generate arbitrary scalars for ext dtype")
117        }
118    })
119}
120
121/// Generates an arbitrary [`PValue`] for the given [`PType`].
122fn random_pvalue(u: &mut Unstructured, ptype: &PType) -> Result<PValue> {
123    Ok(match ptype {
124        PType::U8 => PValue::U8(u.arbitrary()?),
125        PType::U16 => PValue::U16(u.arbitrary()?),
126        PType::U32 => PValue::U32(u.arbitrary()?),
127        PType::U64 => PValue::U64(u.arbitrary()?),
128        PType::I8 => PValue::I8(u.arbitrary()?),
129        PType::I16 => PValue::I16(u.arbitrary()?),
130        PType::I32 => PValue::I32(u.arbitrary()?),
131        PType::I64 => PValue::I64(u.arbitrary()?),
132        PType::F16 => PValue::F16(f16::from_bits(u.arbitrary()?)),
133        PType::F32 => PValue::F32(u.arbitrary()?),
134        PType::F64 => PValue::F64(u.arbitrary()?),
135    })
136}
137
138/// Generates an arbitrary decimal scalar confined to the given bounds of precision and scale.
139///
140/// # Errors
141///
142/// Returns an error if the underlying arbitrary generation fails.
143pub fn random_decimal(u: &mut Unstructured, decimal_type: &DecimalDType) -> Result<ScalarValue> {
144    let precision = decimal_type.precision();
145    let value = match_each_decimal_value_type!(
146        DecimalType::smallest_decimal_value_type(decimal_type),
147        |D| {
148            DecimalValue::from(u.int_in_range(
149                D::MIN_BY_PRECISION[precision as usize]..=D::MAX_BY_PRECISION[precision as usize],
150            )?)
151        }
152    );
153
154    Ok(ScalarValue::Decimal(value))
155}
vortex_array/scalar/arbitrary.rs

vortex_array/scalar/
arbitrary.rs
