sql_cli/benchmarks/
data_generator.rs

1use crate::data::datatable::{DataColumn, DataRow, DataTable, DataType, DataValue};
2use std::collections::HashMap;
3
4pub struct BenchmarkDataGenerator;
5
6impl BenchmarkDataGenerator {
7    pub fn generate_narrow_table(rows: usize) -> DataTable {
8        let mut table = DataTable::new("narrow_bench");
9
10        // Add columns
11        table.add_column(DataColumn::new("id").with_type(DataType::Integer));
12        table.add_column(DataColumn::new("amount").with_type(DataType::Float));
13        table.add_column(DataColumn::new("category_id").with_type(DataType::Integer));
14
15        // Add rows
16        for i in 1..=rows {
17            let row = DataRow::new(vec![
18                DataValue::Integer(i as i64),
19                DataValue::Float((i as f64) * 1.5 + 100.0),
20                DataValue::Integer(((i % 100) + 1) as i64),
21            ]);
22            table.add_row(row).unwrap();
23        }
24
25        table
26    }
27
28    pub fn generate_wide_table(rows: usize, columns: usize) -> DataTable {
29        let mut table = DataTable::new("wide_bench");
30
31        // Add columns
32        table.add_column(DataColumn::new("id").with_type(DataType::Integer));
33        for i in 1..columns {
34            let col_type = match i % 3 {
35                0 => DataType::Integer,
36                1 => DataType::Float,
37                _ => DataType::String,
38            };
39            table.add_column(DataColumn::new(&format!("col_{}", i)).with_type(col_type));
40        }
41
42        // Add rows
43        for i in 1..=rows {
44            let mut row_values = vec![DataValue::Integer(i as i64)];
45
46            for j in 1..columns {
47                let value = match j % 3 {
48                    0 => DataValue::Integer((i * j) as i64),
49                    1 => DataValue::Float((i * j) as f64 * 0.1),
50                    _ => DataValue::String(format!("val_{}_{}", i % 10, j % 5)),
51                };
52                row_values.push(value);
53            }
54
55            table.add_row(DataRow::new(row_values)).unwrap();
56        }
57
58        table
59    }
60
61    pub fn generate_mixed_type_table(rows: usize) -> DataTable {
62        let mut table = DataTable::new("mixed_bench");
63
64        // Add columns
65        table.add_column(DataColumn::new("id").with_type(DataType::Integer));
66        table.add_column(DataColumn::new("category").with_type(DataType::String));
67        table.add_column(DataColumn::new("price").with_type(DataType::Float));
68        table.add_column(DataColumn::new("quantity").with_type(DataType::Integer));
69        table.add_column(DataColumn::new("status").with_type(DataType::String));
70        table.add_column(DataColumn::new("description").with_type(DataType::String));
71
72        let statuses = vec!["Active", "Inactive", "Pending", "Completed", "Failed"];
73        let categories = vec!["Electronics", "Books", "Clothing", "Food", "Toys"];
74
75        // Add rows
76        for i in 1..=rows {
77            let row = DataRow::new(vec![
78                DataValue::Integer(i as i64),
79                DataValue::String(categories[i % 5].to_string()),
80                DataValue::Float(((i * 13) % 1000) as f64 + 10.0),
81                DataValue::Integer(((i * 7) % 100 + 1) as i64),
82                DataValue::String(statuses[i % 5].to_string()),
83                DataValue::String(format!("Product description for item {}", i)),
84            ]);
85            table.add_row(row).unwrap();
86        }
87
88        table
89    }
90
91    pub fn generate_aggregation_table(rows: usize) -> DataTable {
92        let mut table = DataTable::new("aggregation_bench");
93
94        // Add columns
95        table.add_column(DataColumn::new("id").with_type(DataType::Integer));
96        table.add_column(DataColumn::new("group_id").with_type(DataType::Integer));
97        table.add_column(DataColumn::new("sub_group").with_type(DataType::String));
98        table.add_column(DataColumn::new("value1").with_type(DataType::Float));
99        table.add_column(DataColumn::new("value2").with_type(DataType::Float));
100        table.add_column(DataColumn::new("value3").with_type(DataType::Float));
101
102        // Add rows
103        for i in 1..=rows {
104            let row = DataRow::new(vec![
105                DataValue::Integer(i as i64),
106                DataValue::Integer(((i - 1) / 100 + 1) as i64), // Creates groups of ~100
107                DataValue::String(format!("sg_{}", (i % 10))),
108                DataValue::Float((i as f64) * 0.1),
109                DataValue::Float((i as f64) * 0.2 + 50.0),
110                DataValue::Float((i as f64) * 0.3 - 25.0),
111            ]);
112            table.add_row(row).unwrap();
113        }
114
115        table
116    }
117
118    pub fn generate_window_function_table(rows: usize) -> DataTable {
119        let mut table = DataTable::new("window_bench");
120
121        // Add columns
122        table.add_column(DataColumn::new("id").with_type(DataType::Integer));
123        table.add_column(DataColumn::new("timestamp").with_type(DataType::Integer));
124        table.add_column(DataColumn::new("department").with_type(DataType::String));
125        table.add_column(DataColumn::new("employee_id").with_type(DataType::Integer));
126        table.add_column(DataColumn::new("sales").with_type(DataType::Float));
127        table.add_column(DataColumn::new("commission").with_type(DataType::Float));
128
129        let departments = vec!["Sales", "Marketing", "Engineering", "HR", "Finance"];
130
131        // Add rows
132        for i in 1..=rows {
133            let row = DataRow::new(vec![
134                DataValue::Integer(i as i64),
135                DataValue::Integer((i * 3600) as i64), // Simulating timestamps
136                DataValue::String(departments[i % 5].to_string()),
137                DataValue::Integer(((i % 20) + 1) as i64),
138                DataValue::Float(((i * 17) % 1000) as f64 + 100.0),
139                DataValue::Float(((i * 17) % 100) as f64 * 0.01),
140            ]);
141            table.add_row(row).unwrap();
142        }
143
144        table
145    }
146
147    pub fn save_benchmark_data(table: &DataTable, filename: &str) -> Result<(), String> {
148        use std::fs::File;
149        use std::io::Write;
150
151        let csv_content = table.to_csv();
152        let mut file =
153            File::create(filename).map_err(|e| format!("Failed to create file: {}", e))?;
154        file.write_all(csv_content.as_bytes())
155            .map_err(|e| format!("Failed to write file: {}", e))?;
156        Ok(())
157    }
158
159    pub fn generate_all_benchmark_tables(base_rows: usize) -> HashMap<String, DataTable> {
160        let mut tables = HashMap::new();
161
162        tables.insert("narrow".to_string(), Self::generate_narrow_table(base_rows));
163        tables.insert("wide".to_string(), Self::generate_wide_table(base_rows, 20));
164        tables.insert(
165            "very_wide".to_string(),
166            Self::generate_wide_table(base_rows, 50),
167        );
168        tables.insert(
169            "mixed".to_string(),
170            Self::generate_mixed_type_table(base_rows),
171        );
172        tables.insert(
173            "aggregation".to_string(),
174            Self::generate_aggregation_table(base_rows),
175        );
176        tables.insert(
177            "window".to_string(),
178            Self::generate_window_function_table(base_rows),
179        );
180
181        tables
182    }
183}
sql_cli/benchmarks/data_generator.rs

sql_cli/benchmarks/
data_generator.rs
