1use std::collections::HashMap;
7
8use crate::result_map::RowData;
9use crate::value::Value;
10
11#[derive(Debug, Clone)]
13pub struct ColumnarSchema {
14 pub names: Vec<String>,
15 pub types: Vec<String>,
16}
17
18impl ColumnarSchema {
19 pub fn new(names: Vec<String>, types: Vec<String>) -> Self {
20 Self { names, types }
21 }
22
23 pub fn column_count(&self) -> usize {
24 self.names.len()
25 }
26
27 pub fn name_index(&self, name: &str) -> Option<usize> {
28 self.names.iter().position(|n| n == name)
29 }
30}
31
32pub struct ColumnarResultSet {
37 columns: Vec<Vec<Value>>,
38 schema: ColumnarSchema,
39 row_count: usize,
40}
41
42impl ColumnarResultSet {
43 pub fn new(schema: ColumnarSchema) -> Self {
44 let column_count = schema.column_count();
45 Self {
46 columns: vec![Vec::new(); column_count],
47 schema,
48 row_count: 0,
49 }
50 }
51
52 pub fn from_row_data(rows: &[RowData], schema: ColumnarSchema) -> Self {
54 let column_count = schema.column_count();
55 let mut columns = vec![Vec::with_capacity(rows.len()); column_count];
56
57 for row in rows {
58 for (i, name) in schema.names.iter().enumerate() {
59 let value = row.get(name).cloned().unwrap_or(Value::Null);
60 columns[i].push(value);
61 }
62 }
63
64 Self {
65 columns,
66 schema,
67 row_count: rows.len(),
68 }
69 }
70
71 pub fn to_row_data(&self) -> Vec<RowData> {
73 let mut rows = Vec::with_capacity(self.row_count);
74
75 for row_idx in 0..self.row_count {
76 let mut map = HashMap::new();
77 for (col_idx, name) in self.schema.names.iter().enumerate() {
78 if col_idx < self.columns.len() && row_idx < self.columns[col_idx].len() {
79 map.insert(name.clone(), self.columns[col_idx][row_idx].clone());
80 }
81 }
82 rows.push(RowData::new(map));
83 }
84
85 rows
86 }
87
88 pub fn column(&self, name: &str) -> Option<&Vec<Value>> {
90 self.schema.name_index(name).map(|i| &self.columns[i])
91 }
92
93 pub fn column_by_index(&self, idx: usize) -> Option<&Vec<Value>> {
95 self.columns.get(idx)
96 }
97
98 pub fn row_count(&self) -> usize {
100 self.row_count
101 }
102
103 pub fn column_count(&self) -> usize {
105 self.schema.column_count()
106 }
107
108 pub fn schema(&self) -> &ColumnarSchema {
110 &self.schema
111 }
112
113 pub fn get(&self, row_idx: usize, col_name: &str) -> Option<&Value> {
115 let col_idx = self.schema.name_index(col_name)?;
116 self.columns.get(col_idx)?.get(row_idx)
117 }
118}
119
120#[cfg(test)]
123mod tests {
124 use super::*;
125
126 fn make_schema() -> ColumnarSchema {
127 ColumnarSchema::new(
128 vec!["id".into(), "name".into(), "age".into()],
129 vec!["INTEGER".into(), "VARCHAR".into(), "INTEGER".into()],
130 )
131 }
132
133 fn make_rows() -> Vec<RowData> {
134 vec![
135 RowData::new(HashMap::from([
136 ("id".into(), Value::I64(1)),
137 ("name".into(), Value::String("Alice".into())),
138 ("age".into(), Value::I32(30)),
139 ])),
140 RowData::new(HashMap::from([
141 ("id".into(), Value::I64(2)),
142 ("name".into(), Value::String("Bob".into())),
143 ("age".into(), Value::I32(25)),
144 ])),
145 RowData::new(HashMap::from([
146 ("id".into(), Value::I64(3)),
147 ("name".into(), Value::String("Charlie".into())),
148 ("age".into(), Value::I32(35)),
149 ])),
150 ]
151 }
152
153 #[test]
154 fn test_columnar_schema() {
155 let schema = make_schema();
156 assert_eq!(schema.column_count(), 3);
157 assert_eq!(schema.name_index("id"), Some(0));
158 assert_eq!(schema.name_index("name"), Some(1));
159 assert_eq!(schema.name_index("nonexistent"), None);
160 }
161
162 #[test]
163 fn test_columnar_result_set_new() {
164 let schema = make_schema();
165 let result = ColumnarResultSet::new(schema);
166 assert_eq!(result.row_count(), 0);
167 assert_eq!(result.column_count(), 3);
168 }
169
170 #[test]
171 fn test_from_row_data() {
172 let rows = make_rows();
173 let schema = make_schema();
174 let result = ColumnarResultSet::from_row_data(&rows, schema);
175
176 assert_eq!(result.row_count(), 3);
177 assert_eq!(result.column_count(), 3);
178 }
179
180 #[test]
181 fn test_column_access() {
182 let rows = make_rows();
183 let schema = make_schema();
184 let result = ColumnarResultSet::from_row_data(&rows, schema);
185
186 let id_col = result.column("id").expect("id column");
187 assert_eq!(id_col.len(), 3);
188 assert_eq!(id_col[0], Value::I64(1));
189 assert_eq!(id_col[1], Value::I64(2));
190 assert_eq!(id_col[2], Value::I64(3));
191
192 let name_col = result.column("name").expect("name column");
193 assert_eq!(name_col[0], Value::String("Alice".into()));
194
195 assert!(result.column("nonexistent").is_none());
196 }
197
198 #[test]
199 fn test_get_value() {
200 let rows = make_rows();
201 let schema = make_schema();
202 let result = ColumnarResultSet::from_row_data(&rows, schema);
203
204 assert_eq!(result.get(0, "name"), Some(&Value::String("Alice".into())));
205 assert_eq!(result.get(1, "name"), Some(&Value::String("Bob".into())));
206 assert_eq!(result.get(2, "age"), Some(&Value::I32(35)));
207 assert_eq!(result.get(3, "id"), None);
208 assert_eq!(result.get(0, "nonexistent"), None);
209 }
210
211 #[test]
212 fn test_roundtrip_row_to_columnar_to_row() {
213 let rows = make_rows();
214 let schema = make_schema();
215 let result = ColumnarResultSet::from_row_data(&rows, schema);
216 let back = result.to_row_data();
217
218 assert_eq!(back.len(), rows.len());
219
220 for (original, converted) in rows.iter().zip(back.iter()) {
221 for name in ["id", "name", "age"] {
222 assert_eq!(
223 original.get(name).cloned(),
224 converted.get(name).cloned(),
225 "列 {} 往返应一致",
226 name
227 );
228 }
229 }
230 }
231
232 #[test]
233 fn test_empty_rows() {
234 let schema = make_schema();
235 let result = ColumnarResultSet::from_row_data(&[], schema);
236 assert_eq!(result.row_count(), 0);
237 assert_eq!(result.to_row_data().len(), 0);
238 }
239
240 #[test]
241 fn test_column_lengths_equal_row_count() {
242 let rows = make_rows();
243 let schema = make_schema();
244 let result = ColumnarResultSet::from_row_data(&rows, schema);
245
246 for name in ["id", "name", "age"] {
247 let col = result.column(name).expect("column");
248 assert_eq!(
249 col.len(),
250 result.row_count(),
251 "列 {} 长度应等于 row_count",
252 name
253 );
254 }
255 }
256
257 #[test]
258 fn test_column_order_matches_schema() {
259 let rows = make_rows();
260 let schema = make_schema();
261 let result = ColumnarResultSet::from_row_data(&rows, schema);
262
263 let id_col = result.column_by_index(0).expect("index 0");
264 let name_col = result.column_by_index(1).expect("index 1");
265 let age_col = result.column_by_index(2).expect("index 2");
266
267 assert_eq!(id_col[0], Value::I64(1));
268 assert_eq!(name_col[0], Value::String("Alice".into()));
269 assert_eq!(age_col[0], Value::I32(30));
270 }
271
272 #[test]
273 fn test_large_dataset_roundtrip() {
274 let n = 10000;
275 let rows: Vec<RowData> = (0..n)
276 .map(|i| {
277 RowData::new(HashMap::from([
278 ("id".into(), Value::I64(i as i64)),
279 ("name".into(), Value::String(format!("user_{}", i))),
280 ("age".into(), Value::I32((i % 100) as i32 + 20)),
281 ]))
282 })
283 .collect();
284
285 let schema = ColumnarSchema::new(
286 vec!["id".into(), "name".into(), "age".into()],
287 vec!["i64".into(), "string".into(), "i32".into()],
288 );
289 let result = ColumnarResultSet::from_row_data(&rows, schema);
290 assert_eq!(result.row_count(), n);
291
292 let back = result.to_row_data();
293 assert_eq!(back.len(), n);
294
295 for (original, converted) in rows.iter().zip(back.iter()) {
296 assert_eq!(original.get("id"), converted.get("id"));
297 assert_eq!(original.get("name"), converted.get("name"));
298 assert_eq!(original.get("age"), converted.get("age"));
299 }
300 }
301
302 #[test]
303 fn test_columnar_batch_aggregation() {
304 let n = 1000;
305 let rows: Vec<RowData> = (0..n)
306 .map(|i| {
307 RowData::new(HashMap::from([
308 ("id".into(), Value::I64(i as i64)),
309 ("value".into(), Value::I64((i % 10) as i64)),
310 ]))
311 })
312 .collect();
313
314 let schema = ColumnarSchema::new(
315 vec!["id".into(), "value".into()],
316 vec!["i64".into(), "i64".into()],
317 );
318 let result = ColumnarResultSet::from_row_data(&rows, schema);
319
320 let value_col = result.column("value").expect("value column");
321 let sum: i64 = value_col
322 .iter()
323 .filter_map(|v| match v {
324 Value::I64(n) => Some(*n),
325 _ => None,
326 })
327 .sum();
328
329 let expected: i64 = (0..n).map(|i| (i % 10) as i64).sum();
330 assert_eq!(sum, expected);
331 }
332}