1use crate::query::AggregateFunction;
4use crate::value::Value;
5use rkyv::{Archive, Deserialize, Serialize};
6use serde::{Deserialize as SerdeDeserialize, Serialize as SerdeSerialize};
7
8#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
12pub struct EntityBlock {
13 pub entity: String,
15 pub ids: Vec<[u8; 16]>,
17 pub columns: Vec<ColumnData>,
19}
20
21#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
23pub struct ColumnData {
24 pub name: String,
26 pub values: Vec<Value>,
28}
29
30impl ColumnData {
31 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 pub fn new(entity: impl Into<String>) -> Self {
43 Self {
44 entity: entity.into(),
45 ids: vec![],
46 columns: vec![],
47 }
48 }
49
50 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 pub fn len(&self) -> usize {
65 self.ids.len()
66 }
67
68 pub fn is_empty(&self) -> bool {
70 self.ids.is_empty()
71 }
72
73 pub fn column(&self, name: &str) -> Option<&ColumnData> {
75 self.columns.iter().find(|c| c.name == name)
76 }
77
78 pub fn get(&self, row: usize, column: &str) -> Option<&Value> {
80 self.column(column).and_then(|c| c.values.get(row))
81 }
82
83 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#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
98pub struct EdgeBlock {
99 pub relation: String,
101 pub edges: Vec<Edge>,
103}
104
105#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
107pub struct Edge {
108 pub from_id: [u8; 16],
110 pub to_id: [u8; 16],
112}
113
114impl Edge {
115 pub fn new(from_id: [u8; 16], to_id: [u8; 16]) -> Self {
117 Self { from_id, to_id }
118 }
119}
120
121impl EdgeBlock {
122 pub fn new(relation: impl Into<String>) -> Self {
124 Self {
125 relation: relation.into(),
126 edges: vec![],
127 }
128 }
129
130 pub fn with_edges(relation: impl Into<String>, edges: Vec<Edge>) -> Self {
132 Self {
133 relation: relation.into(),
134 edges,
135 }
136 }
137
138 pub fn len(&self) -> usize {
140 self.edges.len()
141 }
142
143 pub fn is_empty(&self) -> bool {
145 self.edges.is_empty()
146 }
147
148 pub fn push(&mut self, edge: Edge) {
150 self.edges.push(edge);
151 }
152}
153
154#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
156pub struct QueryResult {
157 pub entities: Vec<EntityBlock>,
159 pub edges: Vec<EdgeBlock>,
161 pub has_more: bool,
163}
164
165impl QueryResult {
166 pub fn empty() -> Self {
168 Self {
169 entities: vec![],
170 edges: vec![],
171 has_more: false,
172 }
173 }
174
175 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 pub fn entity_block(&self, entity: &str) -> Option<&EntityBlock> {
186 self.entities.iter().find(|b| b.entity == entity)
187 }
188
189 pub fn edge_block(&self, relation: &str) -> Option<&EdgeBlock> {
191 self.edges.iter().find(|b| b.relation == relation)
192 }
193
194 pub fn total_entities(&self) -> usize {
196 self.entities.iter().map(|b| b.len()).sum()
197 }
198
199 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#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
213pub struct MutationResult {
214 pub affected: u64,
216 pub inserted_ids: Vec<[u8; 16]>,
218}
219
220impl MutationResult {
221 pub fn inserted(id: [u8; 16]) -> Self {
223 Self {
224 affected: 1,
225 inserted_ids: vec![id],
226 }
227 }
228
229 pub fn affected(count: u64) -> Self {
231 Self {
232 affected: count,
233 inserted_ids: vec![],
234 }
235 }
236
237 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#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
249pub struct AggregateResult {
250 pub entity: String,
252 pub values: Vec<AggregateValue>,
254}
255
256impl AggregateResult {
257 pub fn new(entity: impl Into<String>, values: Vec<AggregateValue>) -> Self {
259 Self {
260 entity: entity.into(),
261 values,
262 }
263 }
264
265 pub fn empty(entity: impl Into<String>) -> Self {
267 Self {
268 entity: entity.into(),
269 values: vec![],
270 }
271 }
272
273 pub fn first_value(&self) -> Option<&Value> {
275 self.values.first().map(|v| &v.value)
276 }
277
278 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#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize, SerdeSerialize, SerdeDeserialize)]
292pub struct AggregateValue {
293 pub function: AggregateFunction,
295 pub field: Option<String>,
297 pub value: Value,
299}
300
301impl AggregateValue {
302 pub fn new(function: AggregateFunction, field: Option<String>, value: Value) -> Self {
304 Self {
305 function,
306 field,
307 value,
308 }
309 }
310
311 pub fn count(value: i64) -> Self {
313 Self {
314 function: AggregateFunction::Count,
315 field: None,
316 value: Value::Int64(value),
317 }
318 }
319
320 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 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 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 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 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); }
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}