Skip to main content

sz_orm_core/
columnar.rs

1//! v3.2.0 零拷贝序列化 — 列式结果集
2//!
3//! `ColumnarResultSet` 按列连续存储数据(`Vec<Vec<Value>>` 外层每元素为一列),
4//! 相比行式存储(`Vec<RowData>`),列式布局对批量遍历单列场景缓存友好。
5
6use std::collections::HashMap;
7
8use crate::result_map::RowData;
9use crate::value::Value;
10
11/// 列式 schema(列名 + 类型)
12#[derive(Debug, Clone)]
13pub struct ColumnarSchema {
14    /// 列名列表
15    pub names: Vec<String>,
16    /// 列类型列表
17    pub types: Vec<String>,
18}
19
20impl ColumnarSchema {
21    /// 创建新的列式 schema
22    pub fn new(names: Vec<String>, types: Vec<String>) -> Self {
23        Self { names, types }
24    }
25
26    /// 返回列数
27    pub fn column_count(&self) -> usize {
28        self.names.len()
29    }
30
31    /// 按列名查找列索引
32    pub fn name_index(&self, name: &str) -> Option<usize> {
33        self.names.iter().position(|n| n == name)
34    }
35}
36
37/// 列式结果集
38///
39/// 数据按列连续存储:`columns[i]` 是第 i 列的所有行值。
40/// 对单列批量遍历(如聚合、过滤)缓存友好。
41pub struct ColumnarResultSet {
42    columns: Vec<Vec<Value>>,
43    schema: ColumnarSchema,
44    row_count: usize,
45}
46
47impl ColumnarResultSet {
48    /// 创建空的列式结果集
49    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    /// 从行式数据转换为列式
59    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    /// 转换回行式数据
78    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    /// 按列名取列数据
95    pub fn column(&self, name: &str) -> Option<&Vec<Value>> {
96        self.schema.name_index(name).map(|i| &self.columns[i])
97    }
98
99    /// 按列索引取列数据
100    pub fn column_by_index(&self, idx: usize) -> Option<&Vec<Value>> {
101        self.columns.get(idx)
102    }
103
104    /// 行数
105    pub fn row_count(&self) -> usize {
106        self.row_count
107    }
108
109    /// 列数
110    pub fn column_count(&self) -> usize {
111        self.schema.column_count()
112    }
113
114    /// schema 引用
115    pub fn schema(&self) -> &ColumnarSchema {
116        &self.schema
117    }
118
119    /// 获取指定行指定列的值
120    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// ─── 单元测试 ─────────────────────────────────────────────────────
127
128#[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}