vortex_array/arrays/dict/
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 num_traits::NumCast;
8use vortex_buffer::Buffer;
9use vortex_error::VortexExpect;
10
11use super::DictArray;
12use crate::ArrayRef;
13use crate::IntoArray;
14use crate::arrays::PrimitiveArray;
15use crate::arrays::arbitrary::ArbitraryArray;
16use crate::arrays::arbitrary::ArbitraryArrayConfig;
17use crate::arrays::arbitrary::ArbitraryWith;
18use crate::arrays::arbitrary::random_validity;
19use crate::dtype::DType;
20use crate::dtype::NativePType;
21use crate::dtype::Nullability;
22use crate::dtype::PType;
23
24/// A wrapper type to implement `Arbitrary` for `DictArray`.
25#[derive(Clone, Debug)]
26pub struct ArbitraryDictArray(pub DictArray);
27
28impl<'a> Arbitrary<'a> for ArbitraryDictArray {
29    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
30        let dtype: DType = u.arbitrary()?;
31        Self::with_dtype(u, &dtype, None)
32    }
33}
34
35impl ArbitraryDictArray {
36    /// Generate an arbitrary DictArray with the given dtype for values.
37    pub fn with_dtype(u: &mut Unstructured, dtype: &DType, len: Option<usize>) -> Result<Self> {
38        // Generate the number of unique values (dictionary size)
39        let values_len = u.int_in_range(1..=20)?;
40        let values = ArbitraryArray::arbitrary_with_config(
41            u,
42            &ArbitraryArrayConfig {
43                dtype: Some(dtype.clone()),
44                len: values_len..=values_len,
45            },
46        )?
47        .0;
48
49        // Generate codes that index into the values
50        let codes_len = len.unwrap_or(u.int_in_range(0..=100)?);
51
52        // Determine the minimum PType that can represent all indices (max index is values_len - 1)
53        let min_codes_ptype = PType::min_unsigned_ptype_for_value((values_len - 1) as u64);
54
55        // Choose a random PType at least as wide as the minimum
56        let valid_ptypes: &[PType] = match min_codes_ptype {
57            PType::U8 => &[
58                PType::U8,
59                PType::U16,
60                PType::U32,
61                PType::U64,
62                PType::I8,
63                PType::I16,
64                PType::I32,
65                PType::I64,
66            ],
67            PType::U16 => &[
68                PType::U16,
69                PType::U32,
70                PType::U64,
71                PType::I16,
72                PType::I32,
73                PType::I64,
74            ],
75            PType::U32 => &[PType::U32, PType::U64, PType::I32, PType::I64],
76            PType::U64 => &[PType::U64, PType::I64],
77            _ => unreachable!(),
78        };
79        let codes_ptype = *u.choose(valid_ptypes)?;
80
81        // Generate codes with optional nullability
82        let codes_nullable: Nullability = u.arbitrary()?;
83        let codes = match codes_ptype {
84            PType::U8 => random_codes::<u8>(u, codes_len, values_len, codes_nullable)?,
85            PType::U16 => random_codes::<u16>(u, codes_len, values_len, codes_nullable)?,
86            PType::U32 => random_codes::<u32>(u, codes_len, values_len, codes_nullable)?,
87            PType::U64 => random_codes::<u64>(u, codes_len, values_len, codes_nullable)?,
88            PType::I8 => random_codes::<i8>(u, codes_len, values_len, codes_nullable)?,
89            PType::I16 => random_codes::<i16>(u, codes_len, values_len, codes_nullable)?,
90            PType::I32 => random_codes::<i32>(u, codes_len, values_len, codes_nullable)?,
91            PType::I64 => random_codes::<i64>(u, codes_len, values_len, codes_nullable)?,
92            _ => unreachable!(),
93        };
94
95        Ok(ArbitraryDictArray(
96            DictArray::try_new(codes, values)
97                .vortex_expect("DictArray creation should succeed in arbitrary impl"),
98        ))
99    }
100}
101
102/// Generate random codes for a DictArray with a specific integer type.
103fn random_codes<T>(
104    u: &mut Unstructured,
105    len: usize,
106    max_value: usize,
107    nullability: Nullability,
108) -> Result<ArrayRef>
109where
110    T: NativePType + NumCast,
111{
112    let codes: Vec<T> = (0..len)
113        .map(|_| {
114            let idx = u.int_in_range(0..=max_value - 1)?;
115            // max_value is bounded by T::MAX in the caller, so conversion always succeeds
116            Ok(T::from(idx).vortex_expect("value within type bounds"))
117        })
118        .collect::<Result<Vec<_>>>()?;
119    let validity = random_validity(u, nullability, len)?;
120    Ok(PrimitiveArray::new(Buffer::copy_from(codes), validity).into_array())
121}
vortex_array/arrays/dict/arbitrary.rs

vortex_array/arrays/dict/
arbitrary.rs
