risingwave_common/test_utils/
rand_chunk.rs1use risingwave_common::array::BytesArray;
16
17use crate::array::{
18 BoolArray, DataChunk, DateArray, DecimalArray, F32Array, F64Array, I16Array, I32Array,
19 I64Array, Int256Array, IntervalArray, JsonbArray, SerialArray, TimeArray, TimestampArray,
20 TimestamptzArray, Utf8Array,
21};
22use crate::test_utils::rand_array::seed_rand_array_ref;
23use crate::types::DataType;
24use crate::util::schema_check;
25
26pub fn gen_chunk(data_types: &[DataType], size: usize, seed: u64, null_ratio: f64) -> DataChunk {
27 let mut columns = vec![];
28
29 for d in data_types {
30 columns.push(match d {
31 DataType::Boolean => seed_rand_array_ref::<BoolArray>(size, seed, null_ratio),
32 DataType::Int16 => seed_rand_array_ref::<I16Array>(size, seed, null_ratio),
33 DataType::Int32 => seed_rand_array_ref::<I32Array>(size, seed, null_ratio),
34 DataType::Int64 => seed_rand_array_ref::<I64Array>(size, seed, null_ratio),
35 DataType::Float32 => seed_rand_array_ref::<F32Array>(size, seed, null_ratio),
36 DataType::Float64 => seed_rand_array_ref::<F64Array>(size, seed, null_ratio),
37 DataType::Decimal => seed_rand_array_ref::<DecimalArray>(size, seed, null_ratio),
38 DataType::Date => seed_rand_array_ref::<DateArray>(size, seed, null_ratio),
39 DataType::Varchar => seed_rand_array_ref::<Utf8Array>(size, seed, null_ratio),
40 DataType::Time => seed_rand_array_ref::<TimeArray>(size, seed, null_ratio),
41 DataType::Serial => seed_rand_array_ref::<SerialArray>(size, seed, null_ratio),
42 DataType::Timestamp => seed_rand_array_ref::<TimestampArray>(size, seed, null_ratio),
43 DataType::Timestamptz => {
44 seed_rand_array_ref::<TimestamptzArray>(size, seed, null_ratio)
45 }
46 DataType::Interval => seed_rand_array_ref::<IntervalArray>(size, seed, null_ratio),
47 DataType::Int256 => seed_rand_array_ref::<Int256Array>(size, seed, null_ratio),
48 DataType::Bytea => seed_rand_array_ref::<BytesArray>(size, seed, null_ratio),
49 DataType::Jsonb => seed_rand_array_ref::<JsonbArray>(size, seed, null_ratio),
50 DataType::Vector(_) | DataType::Struct(_) | DataType::List(_) | DataType::Map(_) => {
51 todo!()
52 }
53 });
54 }
55 schema_check::schema_check(data_types, &columns).unwrap();
56 DataChunk::new(columns, size)
57}