risingwave_common/test_utils/
rand_array.rs

1// Copyright 2025 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;
20
21use chrono::Datelike;
22use rand::distr::StandardUniform;
23use rand::prelude::Distribution;
24use rand::rngs::SmallRng;
25use rand::{Rng, SeedableRng};
26
27use crate::array::{Array, ArrayBuilder, ArrayRef, ListValue, MapValue, StructValue, VectorVal};
28use crate::types::{
29    DataType, Date, Decimal, Int256, Interval, JsonbVal, MapType, NativeType, Scalar, Serial, Time,
30    Timestamp, Timestamptz,
31};
32
33pub trait RandValue {
34    fn rand_value<R: Rng>(rand: &mut R) -> Self;
35}
36
37impl<T> RandValue for T
38where
39    T: NativeType,
40    StandardUniform: Distribution<T>,
41{
42    fn rand_value<R: Rng>(rand: &mut R) -> Self {
43        rand.random()
44    }
45}
46
47impl RandValue for Box<str> {
48    fn rand_value<R: Rng>(rand: &mut R) -> Self {
49        let len = rand.random_range(1..=10);
50        // `rand.random:::<char>` create a random Unicode scalar value.
51        (0..len)
52            .map(|_| rand.random::<char>())
53            .collect::<String>()
54            .into_boxed_str()
55    }
56}
57
58impl RandValue for Box<[u8]> {
59    fn rand_value<R: Rng>(rand: &mut R) -> Self {
60        let len = rand.random_range(1..=10);
61        (0..len)
62            .map(|_| rand.random::<char>())
63            .collect::<String>()
64            .into_bytes()
65            .into()
66    }
67}
68
69impl RandValue for Decimal {
70    fn rand_value<R: Rng>(rand: &mut R) -> Self {
71        Decimal::try_from((rand.random::<u32>() as f64) + 0.1f64).unwrap()
72    }
73}
74
75impl RandValue for Interval {
76    fn rand_value<R: Rng>(rand: &mut R) -> Self {
77        let months = rand.random_range(0..100);
78        let days = rand.random_range(0..200);
79        let usecs = rand.random_range(0..100_000);
80        Interval::from_month_day_usec(months, days, usecs)
81    }
82}
83
84impl RandValue for Date {
85    fn rand_value<R: Rng>(rand: &mut R) -> Self {
86        let max_day = chrono::NaiveDate::MAX.num_days_from_ce();
87        let min_day = chrono::NaiveDate::MIN.num_days_from_ce();
88        let days = rand.random_range(min_day..=max_day);
89        Date::with_days_since_ce(days).unwrap()
90    }
91}
92
93impl RandValue for Time {
94    fn rand_value<R: Rng>(rand: &mut R) -> Self {
95        let hour = rand.random_range(0..24);
96        let min = rand.random_range(0..60);
97        let sec = rand.random_range(0..60);
98        let nano = rand.random_range(0..1_000_000_000);
99        Time::from_hms_nano_uncheck(hour, min, sec, nano)
100    }
101}
102
103impl RandValue for Timestamp {
104    fn rand_value<R: Rng>(rand: &mut R) -> Self {
105        Timestamp::new(Date::rand_value(rand).0.and_time(Time::rand_value(rand).0))
106    }
107}
108
109impl RandValue for Timestamptz {
110    fn rand_value<R: Rng>(rand: &mut R) -> Self {
111        Timestamptz::from_micros(rand.random())
112    }
113}
114
115impl RandValue for bool {
116    fn rand_value<R: Rng>(rand: &mut R) -> Self {
117        rand.random::<bool>()
118    }
119}
120
121impl RandValue for Serial {
122    fn rand_value<R: Rng>(rand: &mut R) -> Self {
123        // TODO(peng), serial should be in format of RowId
124        i64::rand_value(rand).into()
125    }
126}
127
128impl RandValue for Int256 {
129    fn rand_value<R: Rng>(rand: &mut R) -> Self {
130        let mut bytes = [0u8; 32];
131        rand.fill_bytes(&mut bytes);
132        Int256::from_ne_bytes(bytes)
133    }
134}
135
136impl RandValue for JsonbVal {
137    fn rand_value<R: rand::Rng>(_rand: &mut R) -> Self {
138        JsonbVal::null()
139    }
140}
141
142impl RandValue for StructValue {
143    fn rand_value<R: rand::Rng>(_rand: &mut R) -> Self {
144        StructValue::new(vec![])
145    }
146}
147
148impl RandValue for ListValue {
149    fn rand_value<R: rand::Rng>(rand: &mut R) -> Self {
150        ListValue::from_iter([rand.random::<i16>()])
151    }
152}
153
154impl RandValue for VectorVal {
155    fn rand_value<R: rand::Rng>(_rand: &mut R) -> Self {
156        todo!("VECTOR_PLACEHOLDER")
157    }
158}
159
160impl RandValue for MapValue {
161    fn rand_value<R: Rng>(_rand: &mut R) -> Self {
162        // dummy value
163        MapValue::from_entries(ListValue::empty(&DataType::Struct(
164            MapType::struct_type_for_map(DataType::Varchar, DataType::Varchar),
165        )))
166    }
167}
168
169pub fn rand_array<A, R>(rand: &mut R, size: usize, null_ratio: f64) -> A
170where
171    A: Array,
172    R: Rng,
173    A::OwnedItem: RandValue,
174{
175    let mut builder = A::Builder::new(size);
176    for _ in 0..size {
177        let is_null = rand.random_bool(null_ratio);
178        if is_null {
179            builder.append_null();
180        } else {
181            let value = A::OwnedItem::rand_value(rand);
182            builder.append(Some(value.as_scalar_ref()));
183        }
184    }
185
186    builder.finish()
187}
188
189pub fn seed_rand_array<A>(size: usize, seed: u64, null_ratio: f64) -> A
190where
191    A: Array,
192    A::OwnedItem: RandValue,
193{
194    let mut rand = SmallRng::seed_from_u64(seed);
195    rand_array(&mut rand, size, null_ratio)
196}
197
198pub fn seed_rand_array_ref<A>(size: usize, seed: u64, null_ratio: f64) -> ArrayRef
199where
200    A: Array,
201    A::OwnedItem: RandValue,
202{
203    let array: A = seed_rand_array(size, seed, null_ratio);
204    Arc::new(array.into())
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210    use crate::for_all_variants;
211
212    #[test]
213    fn test_create_array() {
214        macro_rules! gen_rand_array {
215            ($( { $data_type:ident, $variant_name:ident, $suffix_name:ident, $scalar:ty, $scalar_ref:ty, $array:ty, $builder:ty } ),*) => {
216            $(
217                // todo!("VECTOR_PLACEHOLDER")
218                if stringify!($data_type) != "Vector"
219                {
220                    let array = seed_rand_array::<$array>(10, 1024, 0.5);
221                    assert_eq!(10, array.len());
222                }
223            )*
224        };
225    }
226
227        for_all_variants! { gen_rand_array }
228    }
229}
risingwave_common/test_utils/rand_array.rs

risingwave_common/test_utils/
rand_array.rs
