tests_fuzz/generator/
repartition_expr.rs

1// Copyright 2023 Greptime Team
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 derive_builder::Builder;
16use rand::Rng;
17
18use crate::context::TableContextRef;
19use crate::error::{Error, Result};
20use crate::generator::Generator;
21use crate::ir::partition_expr::{SimplePartitions, generate_unique_bound};
22use crate::ir::repartition_expr::{MergePartitionExpr, SplitPartitionExpr};
23
24#[derive(Builder)]
25#[builder(pattern = "owned")]
26pub struct SplitPartitionExprGenerator {
27    table_ctx: TableContextRef,
28}
29
30impl<R: Rng + 'static> Generator<SplitPartitionExpr, R> for SplitPartitionExprGenerator {
31    type Error = Error;
32
33    fn generate(&self, rng: &mut R) -> Result<SplitPartitionExpr> {
34        let table_name = self.table_ctx.name.clone();
35        let partition_def = self
36            .table_ctx
37            .partition
38            .as_ref()
39            .expect("expected partition def");
40        let mut partitions =
41            SimplePartitions::from_exprs(partition_def.columns[0].clone(), &partition_def.exprs)?;
42        let column = self
43            .table_ctx
44            .columns
45            .iter()
46            .find(|c| c.name.value == partitions.column_name.value)
47            .unwrap_or_else(|| {
48                panic!(
49                    "partition column not found: {}, columns: {:?}",
50                    partitions.column_name.value,
51                    self.table_ctx
52                        .columns
53                        .iter()
54                        .map(|c| &c.name.value)
55                        .collect::<Vec<_>>()
56                )
57            });
58        let new_bound = generate_unique_bound(rng, &column.column_type, &partitions.bounds)?;
59        let insert_pos = partitions.insert_bound(new_bound)?;
60
61        let from_expr = partition_def.exprs[insert_pos].clone();
62        let new_exprs = partitions.generate()?;
63        let left = new_exprs[insert_pos].clone();
64        let right = new_exprs[insert_pos + 1].clone();
65        Ok(SplitPartitionExpr {
66            table_name,
67            target: from_expr,
68            into: vec![left, right],
69        })
70    }
71}
72
73#[derive(Builder)]
74#[builder(pattern = "owned")]
75pub struct MergePartitionExprGenerator {
76    table_ctx: TableContextRef,
77}
78
79impl<R: Rng + 'static> Generator<MergePartitionExpr, R> for MergePartitionExprGenerator {
80    type Error = Error;
81
82    fn generate(&self, rng: &mut R) -> Result<MergePartitionExpr> {
83        let table_name = self.table_ctx.name.clone();
84        let partition_def = self
85            .table_ctx
86            .partition
87            .as_ref()
88            .expect("expected partition def");
89        let mut partitions =
90            SimplePartitions::from_exprs(partition_def.columns[0].clone(), &partition_def.exprs)?;
91        let remove_idx = rng.random_range(0..partitions.bounds.len());
92        partitions.remove_bound(remove_idx)?;
93        let left = partition_def.exprs[remove_idx].clone();
94        let right = partition_def.exprs[remove_idx + 1].clone();
95
96        Ok(MergePartitionExpr {
97            table_name,
98            targets: vec![left, right],
99        })
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use std::sync::Arc;
106
107    use rand::{Rng, SeedableRng};
108
109    use crate::context::TableContext;
110    use crate::generator::Generator;
111    use crate::generator::create_expr::CreateTableExprGeneratorBuilder;
112    use crate::generator::repartition_expr::{
113        MergePartitionExprGeneratorBuilder, SplitPartitionExprGeneratorBuilder,
114    };
115    use crate::ir::repartition_expr::RepartitionExpr;
116    use crate::translator::DslTranslator;
117    use crate::translator::mysql::repartition_expr::RepartitionExprTranslator;
118
119    #[test]
120    fn test_split_partition_expr() {
121        common_telemetry::init_default_ut_logging();
122        let mut rng = rand::rng();
123        let seed = rng.random::<u64>();
124        common_telemetry::info!("initializing rng with seed: {seed}");
125        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(seed);
126        let expr = CreateTableExprGeneratorBuilder::default()
127            .columns(10)
128            .partition(10)
129            .if_not_exists(true)
130            .engine("mito2")
131            .build()
132            .unwrap()
133            .generate(&mut rng)
134            .unwrap();
135        let table_ctx = Arc::new(TableContext::from(&expr));
136        SplitPartitionExprGeneratorBuilder::default()
137            .table_ctx(table_ctx)
138            .build()
139            .unwrap()
140            .generate(&mut rng)
141            .unwrap();
142    }
143
144    #[test]
145    fn test_split_partition_expr_deterministic() {
146        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(1234);
147        let expr = CreateTableExprGeneratorBuilder::default()
148            .columns(10)
149            .partition(10)
150            .if_not_exists(true)
151            .engine("mito2")
152            .build()
153            .unwrap()
154            .generate(&mut rng)
155            .unwrap();
156        let table_ctx = Arc::new(TableContext::from(&expr));
157        let expr = SplitPartitionExprGeneratorBuilder::default()
158            .table_ctx(table_ctx)
159            .build()
160            .unwrap()
161            .generate(&mut rng)
162            .unwrap();
163
164        let sql = RepartitionExprTranslator
165            .translate(&RepartitionExpr::Split(expr))
166            .unwrap();
167        let expected = r#"ALTER TABLE quO SPLIT PARTITION (
168  MaGnI >= 1844674407370955160 AND MaGnI < 2767011611056432740
169) INTO (
170  MaGnI >= 1844674407370955160 AND MaGnI < 2290232243991136014,
171  MaGnI >= 2290232243991136014 AND MaGnI < 2767011611056432740
172);"#;
173        assert_eq!(expected, sql);
174    }
175
176    #[test]
177    fn test_merge_partition_expr() {
178        common_telemetry::init_default_ut_logging();
179        let mut rng = rand::rng();
180        let seed = rng.random::<u64>();
181        common_telemetry::info!("initializing rng with seed: {seed}");
182        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(seed);
183        let expr = CreateTableExprGeneratorBuilder::default()
184            .columns(10)
185            .partition(10)
186            .if_not_exists(true)
187            .engine("mito2")
188            .build()
189            .unwrap()
190            .generate(&mut rng)
191            .unwrap();
192        let table_ctx = Arc::new(TableContext::from(&expr));
193        MergePartitionExprGeneratorBuilder::default()
194            .table_ctx(table_ctx)
195            .build()
196            .unwrap()
197            .generate(&mut rng)
198            .unwrap();
199    }
200
201    #[test]
202    fn test_merge_partition_expr_deterministic() {
203        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(1234);
204        let expr = CreateTableExprGeneratorBuilder::default()
205            .columns(10)
206            .partition(10)
207            .if_not_exists(true)
208            .engine("mito2")
209            .build()
210            .unwrap()
211            .generate(&mut rng)
212            .unwrap();
213        let table_ctx = Arc::new(TableContext::from(&expr));
214        let expr = MergePartitionExprGeneratorBuilder::default()
215            .table_ctx(table_ctx)
216            .build()
217            .unwrap()
218            .generate(&mut rng)
219            .unwrap();
220
221        let sql = RepartitionExprTranslator
222            .translate(&RepartitionExpr::Merge(expr))
223            .unwrap();
224        let expected = r#"ALTER TABLE quO MERGE PARTITION (
225  MaGnI >= 3689348814741910320 AND MaGnI < 4611686018427387900,
226  MaGnI >= 4611686018427387900 AND MaGnI < 5534023222112865480
227);"#;
228        assert_eq!(expected, sql);
229    }
230}
tests_fuzz/generator/repartition_expr.rs

tests_fuzz/generator/
repartition_expr.rs
