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::Map(map, _) => Scalar::try_new(
90            dtype.clone(),
91            Some(ScalarValue::Tuple(
92                iter::from_fn(|| {
93                    u.arbitrary().unwrap_or(false).then(|| {
94                        let key = random_scalar(u, &map.key_dtype())?;
95                        let value = random_scalar(u, &map.value_dtype())?;
96                        Ok(Some(ScalarValue::Tuple(vec![
97                            key.into_value(),
98                            value.into_value(),
99                        ])))
100                    })
101                })
102                .collect::<Result<Vec<_>>>()?,
103            )),
104        )
105        .vortex_expect("unable to construct random `Scalar`_"),
106        DType::Struct(sdt, _) => Scalar::try_new(
107            dtype.clone(),
108            Some(ScalarValue::Tuple(
109                sdt.fields()
110                    .map(|d| random_scalar(u, &d).map(|s| s.into_value()))
111                    .collect::<Result<Vec<_>>>()?,
112            )),
113        )
114        .vortex_expect("unable to construct random `Scalar`_"),
115        DType::Union(variants, nullability) => {
116            let child_index = u.choose_index(variants.len())?;
117
118            let child_dtype = variants
119                .variant_by_index(child_index)
120                .vortex_expect("chosen union child index must be valid");
121            let child = random_scalar(u, &child_dtype)?;
122
123            Scalar::union(
124                variants.clone(),
125                variants.child_index_to_tag(child_index),
126                child,
127                *nullability,
128            )
129            .vortex_expect("generated union scalar must be valid")
130        }
131        DType::Variant(_) => todo!(),
132        DType::Extension(..) => {
133            unreachable!("Can't yet generate arbitrary scalars for ext dtype")
134        }
135    })
136}
137
138/// Generates an arbitrary [`PValue`] for the given [`PType`].
139fn random_pvalue(u: &mut Unstructured, ptype: &PType) -> Result<PValue> {
140    Ok(match ptype {
141        PType::U8 => PValue::U8(u.arbitrary()?),
142        PType::U16 => PValue::U16(u.arbitrary()?),
143        PType::U32 => PValue::U32(u.arbitrary()?),
144        PType::U64 => PValue::U64(u.arbitrary()?),
145        PType::I8 => PValue::I8(u.arbitrary()?),
146        PType::I16 => PValue::I16(u.arbitrary()?),
147        PType::I32 => PValue::I32(u.arbitrary()?),
148        PType::I64 => PValue::I64(u.arbitrary()?),
149        PType::F16 => PValue::F16(f16::from_bits(u.arbitrary()?)),
150        PType::F32 => PValue::F32(u.arbitrary()?),
151        PType::F64 => PValue::F64(u.arbitrary()?),
152    })
153}
154
155/// Generates an arbitrary decimal scalar confined to the given bounds of precision and scale.
156///
157/// # Errors
158///
159/// Returns an error if the underlying arbitrary generation fails.
160pub fn random_decimal(u: &mut Unstructured, decimal_type: &DecimalDType) -> Result<ScalarValue> {
161    let precision = decimal_type.precision();
162    let value = match_each_decimal_value_type!(
163        DecimalType::smallest_decimal_value_type(decimal_type),
164        |D| {
165            DecimalValue::from(u.int_in_range(
166                D::MIN_BY_PRECISION[precision as usize]..=D::MAX_BY_PRECISION[precision as usize],
167            )?)
168        }
169    );
170
171    Ok(ScalarValue::Decimal(value))
172}
vortex_array/scalar/arbitrary.rs

vortex_array/scalar/
arbitrary.rs
