Skip to content

Commit b059373

Browse files
committed
add some new fuzz case based on AggregationFuzzer.
1 parent 8657794 commit b059373

1 file changed

Lines changed: 112 additions & 0 deletions

File tree

datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,118 @@ use rand::rngs::StdRng;
4444
use rand::{Rng, SeedableRng};
4545
use tokio::task::JoinSet;
4646

47+
use crate::fuzz_cases::aggregation_fuzzer::{
48+
AggregationFuzzerBuilder, ColumnDescr, DatasetGeneratorConfig,
49+
};
50+
51+
// ========================================================================
52+
// The new aggregation fuzz tests based on [`AggregationFuzzer`]
53+
// ========================================================================
54+
55+
// TODO: write more test case to cover more `group by`s and `aggregation function`s
56+
// TODO: maybe we can use macro to simply the case creating
57+
58+
/// Fuzz test for group by `single int64`
59+
#[tokio::test(flavor = "multi_thread")]
60+
async fn test_group_by_single_int64() {
61+
let builder = AggregationFuzzerBuilder::default();
62+
63+
// Define data generator config
64+
let columns = vec![
65+
ColumnDescr::new("a", DataType::Int64),
66+
ColumnDescr::new("b", DataType::Int64),
67+
ColumnDescr::new("c", DataType::Int64),
68+
];
69+
let sort_keys_set = vec![
70+
vec!["b".to_string()],
71+
vec!["c".to_string(), "b".to_string()],
72+
];
73+
let data_gen_config = DatasetGeneratorConfig {
74+
columns,
75+
rows_num_range: (512, 1024),
76+
sort_keys_set,
77+
};
78+
79+
// Build fuzzer
80+
let fuzzer = builder
81+
.data_gen_config(data_gen_config)
82+
.data_gen_rounds(20)
83+
.sql("SELECT b, sum(a) FROM fuzz_table GROUP BY b")
84+
.table_name("fuzz_table")
85+
.build();
86+
87+
fuzzer.run().await;
88+
}
89+
90+
/// Fuzz test for group by `single string`
91+
#[tokio::test(flavor = "multi_thread")]
92+
async fn test_group_by_single_string() {
93+
let builder = AggregationFuzzerBuilder::default();
94+
95+
// Define data generator config
96+
let columns = vec![
97+
ColumnDescr::new("a", DataType::Int64),
98+
ColumnDescr::new("b", DataType::Utf8),
99+
ColumnDescr::new("c", DataType::Int64),
100+
];
101+
let sort_keys_set = vec![
102+
vec!["b".to_string()],
103+
vec!["c".to_string(), "b".to_string()],
104+
];
105+
let data_gen_config = DatasetGeneratorConfig {
106+
columns,
107+
rows_num_range: (512, 1024),
108+
sort_keys_set,
109+
};
110+
111+
// Build fuzzer
112+
let fuzzer = builder
113+
.data_gen_config(data_gen_config)
114+
.data_gen_rounds(20)
115+
.sql("SELECT b, sum(a) FROM fuzz_table GROUP BY b")
116+
.table_name("fuzz_table")
117+
.build();
118+
119+
fuzzer.run().await;
120+
}
121+
122+
/// Fuzz test for group by `sting + int64`
123+
#[tokio::test(flavor = "multi_thread")]
124+
async fn test_group_by_mixed_string_int64() {
125+
let builder = AggregationFuzzerBuilder::default();
126+
127+
// Define data generator config
128+
let columns = vec![
129+
ColumnDescr::new("a", DataType::Int64),
130+
ColumnDescr::new("b", DataType::Utf8),
131+
ColumnDescr::new("c", DataType::Int64),
132+
ColumnDescr::new("d", DataType::Int32),
133+
];
134+
let sort_keys_set = vec![
135+
vec!["b".to_string(), "c".to_string()],
136+
vec!["d".to_string(), "b".to_string(), "c".to_string()],
137+
];
138+
let data_gen_config = DatasetGeneratorConfig {
139+
columns,
140+
rows_num_range: (512, 1024),
141+
sort_keys_set,
142+
};
143+
144+
// Build fuzzer
145+
let fuzzer = builder
146+
.data_gen_config(data_gen_config)
147+
.data_gen_rounds(20)
148+
.sql("SELECT b, sum(a) FROM fuzz_table GROUP BY b,c")
149+
.table_name("fuzz_table")
150+
.build();
151+
152+
fuzzer.run().await;
153+
}
154+
155+
// ========================================================================
156+
// The old aggregation fuzz tests
157+
// ========================================================================
158+
/// Tracks if this stream is generating input or output
47159
/// Tests that streaming aggregate and batch (non streaming) aggregate produce
48160
/// same results
49161
#[tokio::test(flavor = "multi_thread")]

0 commit comments

Comments
 (0)