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};
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 MapValue {
155    fn rand_value<R: Rng>(_rand: &mut R) -> Self {
156        // dummy value
157        MapValue::from_entries(ListValue::empty(&DataType::Struct(
158            MapType::struct_type_for_map(DataType::Varchar, DataType::Varchar),
159        )))
160    }
161}
162
163pub fn rand_array<A, R>(rand: &mut R, size: usize, null_ratio: f64) -> A
164where
165    A: Array,
166    R: Rng,
167    A::OwnedItem: RandValue,
168{
169    let mut builder = A::Builder::new(size);
170    for _ in 0..size {
171        let is_null = rand.random_bool(null_ratio);
172        if is_null {
173            builder.append_null();
174        } else {
175            let value = A::OwnedItem::rand_value(rand);
176            builder.append(Some(value.as_scalar_ref()));
177        }
178    }
179
180    builder.finish()
181}
182
183pub fn seed_rand_array<A>(size: usize, seed: u64, null_ratio: f64) -> A
184where
185    A: Array,
186    A::OwnedItem: RandValue,
187{
188    let mut rand = SmallRng::seed_from_u64(seed);
189    rand_array(&mut rand, size, null_ratio)
190}
191
192pub fn seed_rand_array_ref<A>(size: usize, seed: u64, null_ratio: f64) -> ArrayRef
193where
194    A: Array,
195    A::OwnedItem: RandValue,
196{
197    let array: A = seed_rand_array(size, seed, null_ratio);
198    Arc::new(array.into())
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204    use crate::for_all_variants;
205
206    #[test]
207    fn test_create_array() {
208        macro_rules! gen_rand_array {
209            ($( { $data_type:ident, $variant_name:ident, $suffix_name:ident, $scalar:ty, $scalar_ref:ty, $array:ty, $builder:ty } ),*) => {
210            $(
211                {
212                    let array = seed_rand_array::<$array>(10, 1024, 0.5);
213                    assert_eq!(10, array.len());
214                }
215            )*
216        };
217    }
218
219        for_all_variants! { gen_rand_array }
220    }
221}
risingwave_common/test_utils/rand_array.rs

risingwave_common/test_utils/
rand_array.rs
