Skip to main content

ormdb_proto/
result.rs

1//! Result types for query responses.
2
3use crate::query::AggregateFunction;
4use crate::value::Value;
5use rkyv::{Archive, Deserialize, Serialize};
6use serde::{Deserialize as SerdeDeserialize, Serialize as SerdeSerialize};
7
8/// A block of entities from a query result.
9///
10/// Uses column-oriented storage for efficient serialization and processing.
11#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
12pub struct EntityBlock {
13    /// Entity type name.
14    pub entity: String,
15    /// Entity IDs (parallel with column values).
16    pub ids: Vec<[u8; 16]>,
17    /// Column data (each column has same length as ids).
18    pub columns: Vec<ColumnData>,
19}
20
21/// Column data within an entity block.
22#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
23pub struct ColumnData {
24    /// Column (field) name.
25    pub name: String,
26    /// Values for each row.
27    pub values: Vec<Value>,
28}
29
30impl ColumnData {
31    /// Create a new column.
32    pub fn new(name: impl Into<String>, values: Vec<Value>) -> Self {
33        Self {
34            name: name.into(),
35            values,
36        }
37    }
38}
39
40impl EntityBlock {
41    /// Create a new empty entity block.
42    pub fn new(entity: impl Into<String>) -> Self {
43        Self {
44            entity: entity.into(),
45            ids: vec![],
46            columns: vec![],
47        }
48    }
49
50    /// Create an entity block with data.
51    pub fn with_data(
52        entity: impl Into<String>,
53        ids: Vec<[u8; 16]>,
54        columns: Vec<ColumnData>,
55    ) -> Self {
56        Self {
57            entity: entity.into(),
58            ids,
59            columns,
60        }
61    }
62
63    /// Get the number of rows in this block.
64    pub fn len(&self) -> usize {
65        self.ids.len()
66    }
67
68    /// Check if this block is empty.
69    pub fn is_empty(&self) -> bool {
70        self.ids.is_empty()
71    }
72
73    /// Get a column by name.
74    pub fn column(&self, name: &str) -> Option<&ColumnData> {
75        self.columns.iter().find(|c| c.name == name)
76    }
77
78    /// Get the value at a specific row and column.
79    pub fn get(&self, row: usize, column: &str) -> Option<&Value> {
80        self.column(column).and_then(|c| c.values.get(row))
81    }
82
83    /// Iterate over rows as (id, field_values) pairs.
84    pub fn rows(&self) -> impl Iterator<Item = (&[u8; 16], Vec<(&str, &Value)>)> {
85        self.ids.iter().enumerate().map(|(i, id)| {
86            let fields: Vec<(&str, &Value)> = self
87                .columns
88                .iter()
89                .map(|col| (col.name.as_str(), &col.values[i]))
90                .collect();
91            (id, fields)
92        })
93    }
94}
95
96/// A block of edges (relationships) from a query result.
97#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
98pub struct EdgeBlock {
99    /// Relation name.
100    pub relation: String,
101    /// Edges in this block.
102    pub edges: Vec<Edge>,
103}
104
105/// A single edge connecting two entities.
106#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
107pub struct Edge {
108    /// ID of the source entity.
109    pub from_id: [u8; 16],
110    /// ID of the target entity.
111    pub to_id: [u8; 16],
112}
113
114impl Edge {
115    /// Create a new edge.
116    pub fn new(from_id: [u8; 16], to_id: [u8; 16]) -> Self {
117        Self { from_id, to_id }
118    }
119}
120
121impl EdgeBlock {
122    /// Create a new empty edge block.
123    pub fn new(relation: impl Into<String>) -> Self {
124        Self {
125            relation: relation.into(),
126            edges: vec![],
127        }
128    }
129
130    /// Create an edge block with edges.
131    pub fn with_edges(relation: impl Into<String>, edges: Vec<Edge>) -> Self {
132        Self {
133            relation: relation.into(),
134            edges,
135        }
136    }
137
138    /// Get the number of edges in this block.
139    pub fn len(&self) -> usize {
140        self.edges.len()
141    }
142
143    /// Check if this block is empty.
144    pub fn is_empty(&self) -> bool {
145        self.edges.is_empty()
146    }
147
148    /// Add an edge to this block.
149    pub fn push(&mut self, edge: Edge) {
150        self.edges.push(edge);
151    }
152}
153
154/// Complete result of a graph query.
155#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
156pub struct QueryResult {
157    /// Entity blocks (one per entity type in the result).
158    pub entities: Vec<EntityBlock>,
159    /// Edge blocks (one per relation in the result).
160    pub edges: Vec<EdgeBlock>,
161    /// Whether there are more results available (pagination).
162    pub has_more: bool,
163}
164
165impl QueryResult {
166    /// Create an empty query result.
167    pub fn empty() -> Self {
168        Self {
169            entities: vec![],
170            edges: vec![],
171            has_more: false,
172        }
173    }
174
175    /// Create a query result with entity and edge blocks.
176    pub fn new(entities: Vec<EntityBlock>, edges: Vec<EdgeBlock>, has_more: bool) -> Self {
177        Self {
178            entities,
179            edges,
180            has_more,
181        }
182    }
183
184    /// Get an entity block by entity type.
185    pub fn entity_block(&self, entity: &str) -> Option<&EntityBlock> {
186        self.entities.iter().find(|b| b.entity == entity)
187    }
188
189    /// Get an edge block by relation name.
190    pub fn edge_block(&self, relation: &str) -> Option<&EdgeBlock> {
191        self.edges.iter().find(|b| b.relation == relation)
192    }
193
194    /// Get total number of entities across all blocks.
195    pub fn total_entities(&self) -> usize {
196        self.entities.iter().map(|b| b.len()).sum()
197    }
198
199    /// Get total number of edges across all blocks.
200    pub fn total_edges(&self) -> usize {
201        self.edges.iter().map(|b| b.len()).sum()
202    }
203}
204
205impl Default for QueryResult {
206    fn default() -> Self {
207        Self::empty()
208    }
209}
210
211/// Result of a mutation operation.
212#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
213pub struct MutationResult {
214    /// Number of entities affected.
215    pub affected: u64,
216    /// IDs of inserted entities (for inserts).
217    pub inserted_ids: Vec<[u8; 16]>,
218}
219
220impl MutationResult {
221    /// Create a result for a successful insert.
222    pub fn inserted(id: [u8; 16]) -> Self {
223        Self {
224            affected: 1,
225            inserted_ids: vec![id],
226        }
227    }
228
229    /// Create a result for a successful update or delete.
230    pub fn affected(count: u64) -> Self {
231        Self {
232            affected: count,
233            inserted_ids: vec![],
234        }
235    }
236
237    /// Create a result for multiple inserts.
238    pub fn bulk_inserted(ids: Vec<[u8; 16]>) -> Self {
239        let affected = ids.len() as u64;
240        Self {
241            affected,
242            inserted_ids: ids,
243        }
244    }
245}
246
247/// Result of an aggregate query.
248#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
249pub struct AggregateResult {
250    /// Entity type that was aggregated.
251    pub entity: String,
252    /// Results for each aggregation.
253    pub values: Vec<AggregateValue>,
254}
255
256impl AggregateResult {
257    /// Create a new aggregate result.
258    pub fn new(entity: impl Into<String>, values: Vec<AggregateValue>) -> Self {
259        Self {
260            entity: entity.into(),
261            values,
262        }
263    }
264
265    /// Create an empty aggregate result.
266    pub fn empty(entity: impl Into<String>) -> Self {
267        Self {
268            entity: entity.into(),
269            values: vec![],
270        }
271    }
272
273    /// Get the first value (for single-aggregation queries).
274    pub fn first_value(&self) -> Option<&Value> {
275        self.values.first().map(|v| &v.value)
276    }
277
278    /// Get the count value if present.
279    pub fn count(&self) -> Option<i64> {
280        self.values
281            .iter()
282            .find(|v| matches!(v.function, AggregateFunction::Count))
283            .and_then(|v| match &v.value {
284                Value::Int64(n) => Some(*n),
285                _ => None,
286            })
287    }
288}
289
290/// A single aggregated value from an aggregate query.
291#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
292pub struct AggregateValue {
293    /// Function that produced this value.
294    pub function: AggregateFunction,
295    /// Field that was aggregated (None for COUNT(*)).
296    pub field: Option<String>,
297    /// The aggregated value.
298    pub value: Value,
299}
300
301impl AggregateValue {
302    /// Create a new aggregate value.
303    pub fn new(function: AggregateFunction, field: Option<String>, value: Value) -> Self {
304        Self {
305            function,
306            field,
307            value,
308        }
309    }
310
311    /// Create a COUNT(*) result.
312    pub fn count(value: i64) -> Self {
313        Self {
314            function: AggregateFunction::Count,
315            field: None,
316            value: Value::Int64(value),
317        }
318    }
319
320    /// Create a COUNT(field) result.
321    pub fn count_field(field: impl Into<String>, value: i64) -> Self {
322        Self {
323            function: AggregateFunction::Count,
324            field: Some(field.into()),
325            value: Value::Int64(value),
326        }
327    }
328
329    /// Create a SUM result.
330    pub fn sum(field: impl Into<String>, value: f64) -> Self {
331        Self {
332            function: AggregateFunction::Sum,
333            field: Some(field.into()),
334            value: Value::Float64(value),
335        }
336    }
337
338    /// Create an AVG result.
339    pub fn avg(field: impl Into<String>, value: f64) -> Self {
340        Self {
341            function: AggregateFunction::Avg,
342            field: Some(field.into()),
343            value: Value::Float64(value),
344        }
345    }
346
347    /// Create a MIN result.
348    pub fn min(field: impl Into<String>, value: Value) -> Self {
349        Self {
350            function: AggregateFunction::Min,
351            field: Some(field.into()),
352            value,
353        }
354    }
355
356    /// Create a MAX result.
357    pub fn max(field: impl Into<String>, value: Value) -> Self {
358        Self {
359            function: AggregateFunction::Max,
360            field: Some(field.into()),
361            value,
362        }
363    }
364}
365
366#[cfg(test)]
367mod tests {
368    use super::*;
369
370    #[test]
371    fn test_entity_block() {
372        let block = EntityBlock::with_data(
373            "User",
374            vec![[1u8; 16], [2u8; 16]],
375            vec![
376                ColumnData::new("name", vec![Value::String("Alice".into()), Value::String("Bob".into())]),
377                ColumnData::new("age", vec![Value::Int32(30), Value::Int32(25)]),
378            ],
379        );
380
381        assert_eq!(block.entity, "User");
382        assert_eq!(block.len(), 2);
383        assert!(!block.is_empty());
384
385        assert_eq!(
386            block.get(0, "name"),
387            Some(&Value::String("Alice".into()))
388        );
389        assert_eq!(block.get(1, "age"), Some(&Value::Int32(25)));
390        assert_eq!(block.get(2, "name"), None); // Out of bounds
391    }
392
393    #[test]
394    fn test_entity_block_rows() {
395        let block = EntityBlock::with_data(
396            "Post",
397            vec![[1u8; 16]],
398            vec![
399                ColumnData::new("title", vec![Value::String("Hello".into())]),
400                ColumnData::new("views", vec![Value::Int64(100)]),
401            ],
402        );
403
404        let rows: Vec<_> = block.rows().collect();
405        assert_eq!(rows.len(), 1);
406        assert_eq!(rows[0].0, &[1u8; 16]);
407        assert_eq!(rows[0].1.len(), 2);
408    }
409
410    #[test]
411    fn test_edge_block() {
412        let mut block = EdgeBlock::new("user_posts");
413        assert!(block.is_empty());
414
415        block.push(Edge::new([1u8; 16], [10u8; 16]));
416        block.push(Edge::new([1u8; 16], [11u8; 16]));
417
418        assert_eq!(block.len(), 2);
419        assert!(!block.is_empty());
420    }
421
422    #[test]
423    fn test_query_result() {
424        let result = QueryResult::new(
425            vec![
426                EntityBlock::with_data("User", vec![[1u8; 16]], vec![]),
427                EntityBlock::with_data("Post", vec![[10u8; 16], [11u8; 16]], vec![]),
428            ],
429            vec![EdgeBlock::with_edges(
430                "user_posts",
431                vec![
432                    Edge::new([1u8; 16], [10u8; 16]),
433                    Edge::new([1u8; 16], [11u8; 16]),
434                ],
435            )],
436            false,
437        );
438
439        assert_eq!(result.total_entities(), 3);
440        assert_eq!(result.total_edges(), 2);
441        assert!(!result.has_more);
442
443        assert!(result.entity_block("User").is_some());
444        assert!(result.entity_block("Comment").is_none());
445        assert!(result.edge_block("user_posts").is_some());
446    }
447
448    #[test]
449    fn test_result_serialization_roundtrip() {
450        let result = QueryResult::new(
451            vec![EntityBlock::with_data(
452                "User",
453                vec![[1u8; 16], [2u8; 16]],
454                vec![
455                    ColumnData::new(
456                        "name",
457                        vec![Value::String("Alice".into()), Value::String("Bob".into())],
458                    ),
459                    ColumnData::new("active", vec![Value::Bool(true), Value::Bool(false)]),
460                ],
461            )],
462            vec![EdgeBlock::with_edges(
463                "friends",
464                vec![Edge::new([1u8; 16], [2u8; 16])],
465            )],
466            true,
467        );
468
469        let bytes = rkyv::to_bytes::<rkyv::rancor::Error>(&result).unwrap();
470        let archived = rkyv::access::<ArchivedQueryResult, rkyv::rancor::Error>(&bytes).unwrap();
471        let deserialized: QueryResult =
472            rkyv::deserialize::<QueryResult, rkyv::rancor::Error>(archived).unwrap();
473
474        assert_eq!(result, deserialized);
475    }
476}