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        // Generate values array with the given dtype
41        let values = ArbitraryArray::arbitrary_with_config(
42            u,
43            &ArbitraryArrayConfig {
44                dtype: Some(dtype.clone()),
45                len: values_len..=values_len,
46            },
47        )?
48        .0;
49
50        // Generate codes that index into the values
51        let codes_len = len.unwrap_or(u.int_in_range(0..=100)?);
52
53        // Determine the minimum PType that can represent all indices (max index is values_len - 1)
54        let min_codes_ptype = PType::min_unsigned_ptype_for_value((values_len - 1) as u64);
55
56        // Choose a random PType at least as wide as the minimum
57        let valid_ptypes: &[PType] = match min_codes_ptype {
58            PType::U8 => &[
59                PType::U8,
60                PType::U16,
61                PType::U32,
62                PType::U64,
63                PType::I8,
64                PType::I16,
65                PType::I32,
66                PType::I64,
67            ],
68            PType::U16 => &[
69                PType::U16,
70                PType::U32,
71                PType::U64,
72                PType::I16,
73                PType::I32,
74                PType::I64,
75            ],
76            PType::U32 => &[PType::U32, PType::U64, PType::I32, PType::I64],
77            PType::U64 => &[PType::U64, PType::I64],
78            _ => unreachable!(),
79        };
80        let codes_ptype = *u.choose(valid_ptypes)?;
81
82        // Generate codes with optional nullability
83        let codes_nullable: Nullability = u.arbitrary()?;
84        let codes = match codes_ptype {
85            PType::U8 => random_codes::<u8>(u, codes_len, values_len, codes_nullable)?,
86            PType::U16 => random_codes::<u16>(u, codes_len, values_len, codes_nullable)?,
87            PType::U32 => random_codes::<u32>(u, codes_len, values_len, codes_nullable)?,
88            PType::U64 => random_codes::<u64>(u, codes_len, values_len, codes_nullable)?,
89            PType::I8 => random_codes::<i8>(u, codes_len, values_len, codes_nullable)?,
90            PType::I16 => random_codes::<i16>(u, codes_len, values_len, codes_nullable)?,
91            PType::I32 => random_codes::<i32>(u, codes_len, values_len, codes_nullable)?,
92            PType::I64 => random_codes::<i64>(u, codes_len, values_len, codes_nullable)?,
93            _ => unreachable!(),
94        };
95
96        Ok(ArbitraryDictArray(
97            DictArray::try_new(codes, values)
98                .vortex_expect("DictArray creation should succeed in arbitrary impl"),
99        ))
100    }
101}
102
103/// Generate random codes for a DictArray with a specific integer type.
104fn random_codes<T>(
105    u: &mut Unstructured,
106    len: usize,
107    max_value: usize,
108    nullability: Nullability,
109) -> Result<ArrayRef>
110where
111    T: NativePType + NumCast,
112{
113    let codes: Vec<T> = (0..len)
114        .map(|_| {
115            let idx = u.int_in_range(0..=max_value - 1)?;
116            // max_value is bounded by T::MAX in the caller, so conversion always succeeds
117            Ok(T::from(idx).vortex_expect("value within type bounds"))
118        })
119        .collect::<Result<Vec<_>>>()?;
120    let validity = random_validity(u, nullability, len)?;
121    Ok(PrimitiveArray::new(Buffer::copy_from(codes), validity).into_array())
122}
vortex_array/arrays/dict/arbitrary.rs

vortex_array/arrays/dict/
arbitrary.rs
