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_dtype::DType;
16use vortex_dtype::DecimalDType;
17use vortex_dtype::NativeDecimalType;
18use vortex_dtype::PType;
19use vortex_dtype::half::f16;
20use vortex_dtype::match_each_decimal_value_type;
21use vortex_error::VortexExpect;
22
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::Struct(sdt, _) => Scalar::try_new(
68            dtype.clone(),
69            Some(ScalarValue::List(
70                sdt.fields()
71                    .map(|d| random_scalar(u, &d).map(|s| s.into_value()))
72                    .collect::<Result<Vec<_>>>()?,
73            )),
74        )
75        .vortex_expect("unable to construct random `Scalar`_"),
76        DType::List(edt, _) => Scalar::try_new(
77            dtype.clone(),
78            Some(ScalarValue::List(
79                iter::from_fn(|| {
80                    // Generate elements with 1/4 probability.
81                    u.arbitrary()
82                        .unwrap_or(false)
83                        .then(|| random_scalar(u, edt).map(|s| s.into_value()))
84                })
85                .collect::<Result<Vec<_>>>()?,
86            )),
87        )
88        .vortex_expect("unable to construct random `Scalar`_"),
89        DType::FixedSizeList(edt, size, _) => Scalar::try_new(
90            dtype.clone(),
91            Some(ScalarValue::List(
92                (0..*size)
93                    .map(|_| random_scalar(u, edt).map(|s| s.into_value()))
94                    .collect::<Result<Vec<_>>>()?,
95            )),
96        )
97        .vortex_expect("unable to construct random `Scalar`_"),
98        DType::Extension(..) => {
99            unreachable!("Can't yet generate arbitrary scalars for ext dtype")
100        }
101    })
102}
103
104/// Generates an arbitrary [`PValue`] for the given [`PType`].
105fn random_pvalue(u: &mut Unstructured, ptype: &PType) -> Result<PValue> {
106    Ok(match ptype {
107        PType::U8 => PValue::U8(u.arbitrary()?),
108        PType::U16 => PValue::U16(u.arbitrary()?),
109        PType::U32 => PValue::U32(u.arbitrary()?),
110        PType::U64 => PValue::U64(u.arbitrary()?),
111        PType::I8 => PValue::I8(u.arbitrary()?),
112        PType::I16 => PValue::I16(u.arbitrary()?),
113        PType::I32 => PValue::I32(u.arbitrary()?),
114        PType::I64 => PValue::I64(u.arbitrary()?),
115        PType::F16 => PValue::F16(f16::from_bits(u.arbitrary()?)),
116        PType::F32 => PValue::F32(u.arbitrary()?),
117        PType::F64 => PValue::F64(u.arbitrary()?),
118    })
119}
120
121/// Generates an arbitrary decimal scalar confined to the given bounds of precision and scale.
122///
123/// # Errors
124///
125/// Returns an error if the underlying arbitrary generation fails.
126pub fn random_decimal(u: &mut Unstructured, decimal_type: &DecimalDType) -> Result<ScalarValue> {
127    let precision = decimal_type.precision();
128    let value = match_each_decimal_value_type!(
129        DecimalType::smallest_decimal_value_type(decimal_type),
130        |D| {
131            DecimalValue::from(u.int_in_range(
132                D::MIN_BY_PRECISION[precision as usize]..=D::MAX_BY_PRECISION[precision as usize],
133            )?)
134        }
135    );
136
137    Ok(ScalarValue::Decimal(value))
138}
vortex_array/scalar/arbitrary.rs

vortex_array/scalar/
arbitrary.rs
