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

vortex_array/arrays/dict/
arbitrary.rs
