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

vortex_array/scalar/
arbitrary.rs
