risingwave_common/field_generator/
varchar.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
15use rand::distr::Alphanumeric;
16use rand::rngs::StdRng;
17use rand::{Rng, SeedableRng};
18use serde_json::{Value, json};
19
20use super::DEFAULT_LENGTH;
21use crate::types::{Datum, Scalar};
22
23pub struct VarcharRandomVariableLengthField {
24    seed: u64,
25}
26
27impl VarcharRandomVariableLengthField {
28    pub fn new(seed: u64) -> Self {
29        Self { seed }
30    }
31
32    pub fn generate_string(&mut self, offset: u64) -> String {
33        let len = rand::rng().random_range(0..=DEFAULT_LENGTH * 2);
34        StdRng::seed_from_u64(offset ^ self.seed)
35            .sample_iter(&Alphanumeric)
36            .take(len)
37            .map(char::from)
38            .collect()
39    }
40
41    pub fn generate(&mut self, offset: u64) -> Value {
42        json!(self.generate_string(offset))
43    }
44
45    pub fn generate_datum(&mut self, offset: u64) -> Datum {
46        let s = self.generate_string(offset);
47        Some(s.into_boxed_str().to_scalar_value())
48    }
49}
50
51pub struct VarcharRandomFixedLengthField {
52    length: usize,
53    seed: u64,
54}
55
56impl VarcharRandomFixedLengthField {
57    pub fn new(length_option: &Option<usize>, seed: u64) -> Self {
58        let length = if let Some(length) = length_option {
59            *length
60        } else {
61            DEFAULT_LENGTH
62        };
63        Self { length, seed }
64    }
65
66    pub fn generate(&mut self, offset: u64) -> Value {
67        let s: String = StdRng::seed_from_u64(offset ^ self.seed)
68            .sample_iter(&Alphanumeric)
69            .take(self.length)
70            .map(char::from)
71            .collect();
72        json!(s)
73    }
74
75    pub fn generate_datum(&mut self, offset: u64) -> Datum {
76        let s: String = StdRng::seed_from_u64(offset ^ self.seed)
77            .sample_iter(&Alphanumeric)
78            .take(self.length)
79            .map(char::from)
80            .collect();
81        Some(s.into_boxed_str().to_scalar_value())
82    }
83}
84
85pub struct VarcharConstant {}
86
87impl VarcharConstant {
88    const CONSTANT_STRING: &'static str = "2022-03-03";
89
90    pub fn generate_json() -> Value {
91        json!(Self::CONSTANT_STRING)
92    }
93
94    pub fn generate_datum() -> Datum {
95        Some(
96            Self::CONSTANT_STRING
97                .to_owned()
98                .into_boxed_str()
99                .to_scalar_value(),
100        )
101    }
102}