risingwave_common/test_utils/
rand_array.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Contains helper methods for generating random arrays in tests.
16//!
17//! Use [`seed_rand_array`] to generate an random array.
18
19use std::sync::{Arc, LazyLock};
20
21use chrono::Datelike;
22use parking_lot::Mutex;
23use rand::distr::StandardUniform;
24use rand::prelude::{Distribution, StdRng};
25use rand::rngs::SmallRng;
26use rand::{Rng, SeedableRng};
27
28use crate::array::{Array, ArrayBuilder, ArrayRef, VectorVal};
29use crate::types::{
30    Date, Decimal, Int256, Interval, JsonbVal, NativeType, Scalar, Serial, Time, Timestamp,
31    Timestamptz, VariantVal,
32};
33
34pub fn gen_vector_for_test(d: usize) -> VectorVal {
35    static RNG: LazyLock<Mutex<StdRng>> = LazyLock::new(|| Mutex::new(StdRng::seed_from_u64(233)));
36    VectorVal::from_iter((0..d).map(|_| RNG.lock().random::<f32>().try_into().unwrap()))
37}
38
39pub trait RandValue {
40    fn rand_value<R: Rng>(rand: &mut R) -> Self;
41}
42
43impl<T> RandValue for T
44where
45    T: NativeType,
46    StandardUniform: Distribution<T>,
47{
48    fn rand_value<R: Rng>(rand: &mut R) -> Self {
49        rand.random()
50    }
51}
52
53impl RandValue for Box<str> {
54    fn rand_value<R: Rng>(rand: &mut R) -> Self {
55        let len = rand.random_range(1..=10);
56        // `rand.random:::<char>` create a random Unicode scalar value.
57        (0..len)
58            .map(|_| rand.random::<char>())
59            .collect::<String>()
60            .into_boxed_str()
61    }
62}
63
64impl RandValue for Box<[u8]> {
65    fn rand_value<R: Rng>(rand: &mut R) -> Self {
66        let len = rand.random_range(1..=10);
67        (0..len)
68            .map(|_| rand.random::<char>())
69            .collect::<String>()
70            .into_bytes()
71            .into()
72    }
73}
74
75impl RandValue for Decimal {
76    fn rand_value<R: Rng>(rand: &mut R) -> Self {
77        Decimal::try_from((rand.random::<u32>() as f64) + 0.1f64).unwrap()
78    }
79}
80
81impl RandValue for Interval {
82    fn rand_value<R: Rng>(rand: &mut R) -> Self {
83        let months = rand.random_range(0..100);
84        let days = rand.random_range(0..200);
85        let usecs = rand.random_range(0..100_000);
86        Interval::from_month_day_usec(months, days, usecs)
87    }
88}
89
90impl RandValue for Date {
91    fn rand_value<R: Rng>(rand: &mut R) -> Self {
92        let max_day = chrono::NaiveDate::MAX.num_days_from_ce();
93        let min_day = chrono::NaiveDate::MIN.num_days_from_ce();
94        let days = rand.random_range(min_day..=max_day);
95        Date::with_days_since_ce(days).unwrap()
96    }
97}
98
99impl RandValue for Time {
100    fn rand_value<R: Rng>(rand: &mut R) -> Self {
101        let hour = rand.random_range(0..24);
102        let min = rand.random_range(0..60);
103        let sec = rand.random_range(0..60);
104        let nano = rand.random_range(0..1_000_000_000);
105        Time::from_hms_nano_uncheck(hour, min, sec, nano)
106    }
107}
108
109impl RandValue for Timestamp {
110    fn rand_value<R: Rng>(rand: &mut R) -> Self {
111        Timestamp::new(Date::rand_value(rand).0.and_time(Time::rand_value(rand).0))
112    }
113}
114
115impl RandValue for Timestamptz {
116    fn rand_value<R: Rng>(rand: &mut R) -> Self {
117        Timestamp::rand_value(rand).0.and_utc().into()
118    }
119}
120
121impl RandValue for bool {
122    fn rand_value<R: Rng>(rand: &mut R) -> Self {
123        rand.random::<bool>()
124    }
125}
126
127impl RandValue for Serial {
128    fn rand_value<R: Rng>(rand: &mut R) -> Self {
129        // TODO(peng), serial should be in format of RowId
130        i64::rand_value(rand).into()
131    }
132}
133
134impl RandValue for Int256 {
135    fn rand_value<R: Rng>(rand: &mut R) -> Self {
136        let mut bytes = [0u8; 32];
137        rand.fill_bytes(&mut bytes);
138        Int256::from_ne_bytes(bytes)
139    }
140}
141
142impl RandValue for JsonbVal {
143    fn rand_value<R: rand::Rng>(_rand: &mut R) -> Self {
144        JsonbVal::null()
145    }
146}
147
148impl RandValue for VariantVal {
149    fn rand_value<R: rand::Rng>(_rand: &mut R) -> Self {
150        VariantVal::null()
151    }
152}
153
154#[cfg(test)]
155impl RandValue for crate::types::StructValue {
156    fn rand_value<R: rand::Rng>(_rand: &mut R) -> Self {
157        crate::types::StructValue::new(vec![])
158    }
159}
160
161#[cfg(test)]
162impl RandValue for crate::types::ListValue {
163    fn rand_value<R: rand::Rng>(rand: &mut R) -> Self {
164        crate::types::ListValue::from_iter([rand.random::<i16>()])
165    }
166}
167
168#[cfg(test)]
169impl RandValue for crate::types::VectorVal {
170    fn rand_value<R: rand::Rng>(rand: &mut R) -> Self {
171        Self::from_iter(
172            [(); Self::TEST_VECTOR_DIMENSION].map(|()| rand.random::<f32>().try_into().unwrap()),
173        )
174    }
175}
176
177#[cfg(test)]
178impl RandValue for crate::types::MapValue {
179    fn rand_value<R: Rng>(_rand: &mut R) -> Self {
180        use crate::types::DataType;
181        // dummy value
182        crate::types::MapValue::from_entries(crate::types::ListValue::empty(&DataType::Struct(
183            crate::types::MapType::struct_type_for_map(DataType::Varchar, DataType::Varchar),
184        )))
185    }
186}
187
188pub fn rand_array<A, R>(rand: &mut R, size: usize, null_ratio: f64) -> A
189where
190    A: Array,
191    R: Rng,
192    A::OwnedItem: RandValue,
193{
194    let mut builder = A::Builder::new(size);
195    for _ in 0..size {
196        let is_null = rand.random_bool(null_ratio);
197        if is_null {
198            builder.append_null();
199        } else {
200            let value = A::OwnedItem::rand_value(rand);
201            builder.append(Some(value.as_scalar_ref()));
202        }
203    }
204
205    builder.finish()
206}
207
208pub fn seed_rand_array<A>(size: usize, seed: u64, null_ratio: f64) -> A
209where
210    A: Array,
211    A::OwnedItem: RandValue,
212{
213    let mut rand = SmallRng::seed_from_u64(seed);
214    rand_array(&mut rand, size, null_ratio)
215}
216
217pub fn seed_rand_array_ref<A>(size: usize, seed: u64, null_ratio: f64) -> ArrayRef
218where
219    A: Array,
220    A::OwnedItem: RandValue,
221{
222    let array: A = seed_rand_array(size, seed, null_ratio);
223    Arc::new(array.into())
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229    use crate::for_all_variants;
230
231    #[test]
232    fn test_create_array() {
233        macro_rules! gen_rand_array {
234            ($( { $data_type:ident, $variant_name:ident, $suffix_name:ident, $scalar:ty, $scalar_ref:ty, $array:ty, $builder:ty } ),*) => {
235            $(
236
237                let array = seed_rand_array::<$array>(10, 1024, 0.5);
238                assert_eq!(10, array.len());
239            )*
240        };
241    }
242
243        for_all_variants! { gen_rand_array }
244    }
245}
risingwave_common/test_utils/rand_array.rs

risingwave_common/test_utils/
rand_array.rs
