risingwave_common/test_utils/
rand_array.rs1use 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,
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 (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 Timestamptz::from_micros(rand.random())
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 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
148#[cfg(test)]
149impl RandValue for crate::types::StructValue {
150 fn rand_value<R: rand::Rng>(_rand: &mut R) -> Self {
151 crate::types::StructValue::new(vec![])
152 }
153}
154
155#[cfg(test)]
156impl RandValue for crate::types::ListValue {
157 fn rand_value<R: rand::Rng>(rand: &mut R) -> Self {
158 crate::types::ListValue::from_iter([rand.random::<i16>()])
159 }
160}
161
162#[cfg(test)]
163impl RandValue for crate::types::VectorVal {
164 fn rand_value<R: rand::Rng>(rand: &mut R) -> Self {
165 Self::from_iter(
166 [(); Self::TEST_VECTOR_DIMENSION].map(|()| rand.random::<f32>().try_into().unwrap()),
167 )
168 }
169}
170
171#[cfg(test)]
172impl RandValue for crate::types::MapValue {
173 fn rand_value<R: Rng>(_rand: &mut R) -> Self {
174 use crate::types::DataType;
175 crate::types::MapValue::from_entries(crate::types::ListValue::empty(&DataType::Struct(
177 crate::types::MapType::struct_type_for_map(DataType::Varchar, DataType::Varchar),
178 )))
179 }
180}
181
182pub fn rand_array<A, R>(rand: &mut R, size: usize, null_ratio: f64) -> A
183where
184 A: Array,
185 R: Rng,
186 A::OwnedItem: RandValue,
187{
188 let mut builder = A::Builder::new(size);
189 for _ in 0..size {
190 let is_null = rand.random_bool(null_ratio);
191 if is_null {
192 builder.append_null();
193 } else {
194 let value = A::OwnedItem::rand_value(rand);
195 builder.append(Some(value.as_scalar_ref()));
196 }
197 }
198
199 builder.finish()
200}
201
202pub fn seed_rand_array<A>(size: usize, seed: u64, null_ratio: f64) -> A
203where
204 A: Array,
205 A::OwnedItem: RandValue,
206{
207 let mut rand = SmallRng::seed_from_u64(seed);
208 rand_array(&mut rand, size, null_ratio)
209}
210
211pub fn seed_rand_array_ref<A>(size: usize, seed: u64, null_ratio: f64) -> ArrayRef
212where
213 A: Array,
214 A::OwnedItem: RandValue,
215{
216 let array: A = seed_rand_array(size, seed, null_ratio);
217 Arc::new(array.into())
218}
219
220#[cfg(test)]
221mod tests {
222 use super::*;
223 use crate::for_all_variants;
224
225 #[test]
226 fn test_create_array() {
227 macro_rules! gen_rand_array {
228 ($( { $data_type:ident, $variant_name:ident, $suffix_name:ident, $scalar:ty, $scalar_ref:ty, $array:ty, $builder:ty } ),*) => {
229 $(
230
231 let array = seed_rand_array::<$array>(10, 1024, 0.5);
232 assert_eq!(10, array.len());
233 )*
234 };
235 }
236
237 for_all_variants! { gen_rand_array }
238 }
239}