Skip to main content

vortex_array/arrays/constant/
arbitrary.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use arbitrary::Arbitrary;
5use arbitrary::Result;
6use arbitrary::Unstructured;
7use vortex_dtype::DType;
8
9use super::ConstantArray;
10use crate::scalar::arbitrary::random_scalar;
11
12/// A wrapper type to implement `Arbitrary` for `ConstantArray`.
13#[derive(Clone, Debug)]
14pub struct ArbitraryConstantArray(pub ConstantArray);
15
16impl<'a> Arbitrary<'a> for ArbitraryConstantArray {
17    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
18        let dtype: DType = u.arbitrary()?;
19        Self::with_dtype(u, &dtype, None)
20    }
21}
22
23impl ArbitraryConstantArray {
24    /// Generate an arbitrary ConstantArray with the given dtype.
25    pub fn with_dtype(u: &mut Unstructured, dtype: &DType, len: Option<usize>) -> Result<Self> {
26        let scalar = random_scalar(u, dtype)?;
27        let len = len.unwrap_or(u.int_in_range(0..=100)?);
28        Ok(ArbitraryConstantArray(ConstantArray::new(scalar, len)))
29    }
30}