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    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
32/// 列式结果集
33///
34/// 数据按列连续存储:`columns[i]` 是第 i 列的所有行值。
35/// 对单列批量遍历(如聚合、过滤)缓存友好。
36pub 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    /// 从行式数据转换为列式
53    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    /// 转换回行式数据
72    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    /// 按列名取列数据
89    pub fn column(&self, name: &str) -> Option<&Vec<Value>> {
90        self.schema.name_index(name).map(|i| &self.columns[i])
91    }
92
93    /// 按列索引取列数据
94    pub fn column_by_index(&self, idx: usize) -> Option<&Vec<Value>> {
95        self.columns.get(idx)
96    }
97
98    /// 行数
99    pub fn row_count(&self) -> usize {
100        self.row_count
101    }
102
103    /// 列数
104    pub fn column_count(&self) -> usize {
105        self.schema.column_count()
106    }
107
108    /// schema 引用
109    pub fn schema(&self) -> &ColumnarSchema {
110        &self.schema
111    }
112
113    /// 获取指定行指定列的值
114    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// ─── 单元测试 ─────────────────────────────────────────────────────
121
122#[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}