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