1use super::{AttributeMap, EdgeIndex, GraphRecord, Group, NodeIndex};
2use crate::{
3 errors::SchemaError,
4 graphrecord::{GraphRecordAttribute, datatypes::DataType},
5};
6use graphrecords_utils::aliases::GrHashMap;
7#[cfg(feature = "serde")]
8use serde::{Deserialize, Serialize};
9use std::{
10 borrow::Borrow,
11 collections::{HashMap, hash_map::Entry},
12 ops::Deref,
13};
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
17pub enum AttributeType {
18 Categorical,
19 Continuous,
20 Temporal,
21 Unstructured,
22}
23
24impl AttributeType {
25 #[must_use]
26 pub fn infer(data_type: &DataType) -> Self {
27 match data_type {
28 DataType::String | DataType::Null | DataType::Any => Self::Unstructured,
29 DataType::Int | DataType::Float => Self::Continuous,
30 DataType::Bool => Self::Categorical,
31 DataType::DateTime | DataType::Duration => Self::Temporal,
32 DataType::Union((first_dataype, second_dataype)) => {
33 Self::infer(first_dataype).merge(Self::infer(second_dataype))
34 }
35 DataType::Option(dataype) => Self::infer(dataype),
36 }
37 }
38
39 const fn merge(self, other: Self) -> Self {
40 match (self, other) {
41 (Self::Categorical, Self::Unstructured) | (Self::Unstructured, Self::Categorical) => {
42 Self::Unstructured
43 }
44 (Self::Categorical, _) | (_, Self::Categorical) => Self::Categorical,
45 (Self::Continuous, Self::Continuous) => Self::Continuous,
46 (Self::Temporal, Self::Temporal) => Self::Temporal,
47 _ => Self::Unstructured,
48 }
49 }
50}
51
52impl DataType {
53 fn merge(&self, other: &Self) -> Self {
54 if self.evaluate(other) {
55 self.clone()
56 } else {
57 match (self, other) {
58 (Self::Null, _) => Self::Option(Box::new(other.clone())),
59 (_, Self::Null) => Self::Option(Box::new(self.clone())),
60 (_, Self::Any) => Self::Any,
61 (Self::Option(option1), Self::Option(option2)) => {
62 Self::Option(Box::new(option1.merge(option2)))
63 }
64 (Self::Option(option), _) => Self::Option(Box::new(option.merge(other))),
65 (_, Self::Option(option)) => Self::Option(Box::new(self.merge(option))),
66 _ => Self::Union((Box::new(self.clone()), Box::new(other.clone()))),
67 }
68 }
69 }
70}
71
72#[derive(Debug, Clone, PartialEq, Eq)]
73#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
74pub struct AttributeDataType {
75 data_type: DataType,
76 attribute_type: AttributeType,
77}
78
79impl AttributeDataType {
80 fn validate(data_type: &DataType, attribute_type: AttributeType) -> Result<(), SchemaError> {
81 match (attribute_type, data_type) {
82 (AttributeType::Categorical | AttributeType::Unstructured, _)
83 | (AttributeType::Continuous, DataType::Int | DataType::Float | DataType::Null)
84 | (AttributeType::Temporal, DataType::DateTime | DataType::Duration | DataType::Null) => {
85 Ok(())
86 }
87
88 (_, DataType::Option(option)) => Self::validate(option, attribute_type),
89 (_, DataType::Union((first_datatype, second_datatype))) => {
90 Self::validate(first_datatype, attribute_type)?;
91 Self::validate(second_datatype, attribute_type)
92 }
93
94 (AttributeType::Continuous, _) => Err(SchemaError::ContinuousAttributeNotNumeric),
95
96 (AttributeType::Temporal, _) => Err(SchemaError::TemporalAttributeNotTemporal),
97 }
98 }
99
100 pub fn new(data_type: DataType, attribute_type: AttributeType) -> Result<Self, SchemaError> {
101 Self::validate(&data_type, attribute_type)?;
102
103 Ok(Self {
104 data_type,
105 attribute_type,
106 })
107 }
108
109 #[must_use]
110 pub const fn data_type(&self) -> &DataType {
111 &self.data_type
112 }
113
114 #[must_use]
115 pub const fn attribute_type(&self) -> &AttributeType {
116 &self.attribute_type
117 }
118
119 fn merge(&mut self, other: &Self) {
120 match (self.data_type.clone(), other.data_type.clone()) {
121 (DataType::Null, _) => {
122 self.data_type = self.data_type.merge(&other.data_type);
123 self.attribute_type = other.attribute_type;
124 }
125 (_, DataType::Null) => {
126 self.data_type = self.data_type.merge(&other.data_type);
127 }
128 _ => {
129 self.data_type = self.data_type.merge(&other.data_type);
130 self.attribute_type = self.attribute_type.merge(other.attribute_type);
131 }
132 }
133 }
134}
135
136impl From<DataType> for AttributeDataType {
137 fn from(value: DataType) -> Self {
138 let attribute_type = AttributeType::infer(&value);
139
140 Self {
141 data_type: value,
142 attribute_type,
143 }
144 }
145}
146
147impl From<(DataType, AttributeType)> for AttributeDataType {
148 fn from(value: (DataType, AttributeType)) -> Self {
149 Self {
150 data_type: value.0,
151 attribute_type: value.1,
152 }
153 }
154}
155
156#[derive(Debug, Clone, Copy)]
157enum AttributeSchemaKind<'a> {
158 Node(&'a NodeIndex),
159 Edge(&'a EdgeIndex),
160}
161
162impl AttributeSchemaKind<'_> {
163 fn attribute_missing_error(
164 &self,
165 attribute: &GraphRecordAttribute,
166 data_type: &DataType,
167 ) -> SchemaError {
168 match self {
169 Self::Node(node_index) => SchemaError::NodeAttributeMissing {
170 node_index: (*node_index).clone(),
171 attribute: attribute.clone(),
172 data_type: data_type.clone(),
173 },
174 Self::Edge(edge_index) => SchemaError::EdgeAttributeMissing {
175 edge_index: **edge_index,
176 attribute: attribute.clone(),
177 data_type: data_type.clone(),
178 },
179 }
180 }
181
182 fn data_type_mismatch_error(
183 &self,
184 attribute: &GraphRecordAttribute,
185 data_type: &DataType,
186 expected_data_type: &DataType,
187 ) -> SchemaError {
188 match self {
189 Self::Node(node_index) => SchemaError::NodeAttributeDataTypeMismatch {
190 node_index: (*node_index).clone(),
191 attribute: attribute.clone(),
192 data_type: data_type.clone(),
193 expected_data_type: expected_data_type.clone(),
194 },
195 Self::Edge(edge_index) => SchemaError::EdgeAttributeDataTypeMismatch {
196 edge_index: **edge_index,
197 attribute: attribute.clone(),
198 data_type: data_type.clone(),
199 expected_data_type: expected_data_type.clone(),
200 },
201 }
202 }
203
204 fn attributes_not_in_schema_error(&self, attributes: Vec<GraphRecordAttribute>) -> SchemaError {
205 match self {
206 Self::Node(node_index) => SchemaError::NodeAttributesNotInSchema {
207 node_index: (*node_index).clone(),
208 attributes,
209 },
210 Self::Edge(edge_index) => SchemaError::EdgeAttributesNotInSchema {
211 edge_index: **edge_index,
212 attributes,
213 },
214 }
215 }
216}
217
218type AttributeSchemaMapping = HashMap<GraphRecordAttribute, AttributeDataType>;
219
220#[derive(Debug, Clone, PartialEq, Eq, Default)]
221#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
222pub struct AttributeSchema(AttributeSchemaMapping);
223
224impl Deref for AttributeSchema {
225 type Target = AttributeSchemaMapping;
226
227 fn deref(&self) -> &Self::Target {
228 &self.0
229 }
230}
231
232impl<T> From<T> for AttributeSchema
233where
234 T: Into<AttributeSchemaMapping>,
235{
236 fn from(value: T) -> Self {
237 Self(value.into())
238 }
239}
240
241impl AttributeSchema {
242 #[must_use]
243 pub const fn new(mapping: HashMap<GraphRecordAttribute, AttributeDataType>) -> Self {
244 Self(mapping)
245 }
246
247 fn validate(
248 &self,
249 attributes: &AttributeMap,
250 kind: &AttributeSchemaKind,
251 ) -> Result<(), SchemaError> {
252 let mut matched_count = 0;
253 let mut attributes_not_in_schema = Vec::new();
254
255 for (key, value) in attributes {
256 match self.0.get(key) {
257 Some(schema) => {
258 let data_type = DataType::from(value);
259
260 if !schema.data_type.evaluate(&data_type) {
261 return Err(kind.data_type_mismatch_error(
262 key,
263 &data_type,
264 &schema.data_type,
265 ));
266 }
267
268 matched_count += 1;
269 }
270 None => {
271 attributes_not_in_schema.push(key.clone());
272 }
273 }
274 }
275
276 if matched_count < self.0.len() {
277 for (key, schema) in &self.0 {
278 if !attributes.contains_key(key) && !matches!(schema.data_type, DataType::Option(_))
279 {
280 return Err(kind.attribute_missing_error(key, &schema.data_type));
281 }
282 }
283 }
284
285 if !attributes_not_in_schema.is_empty() {
286 return Err(kind.attributes_not_in_schema_error(attributes_not_in_schema));
287 }
288
289 Ok(())
290 }
291
292 fn update(&mut self, attributes: &AttributeMap, empty: bool) {
293 for (attribute, data_type) in &mut self.0 {
294 if !attributes.contains_key(attribute) {
295 data_type.data_type = data_type.data_type.merge(&DataType::Null);
296 }
297 }
298
299 for (attribute, value) in attributes {
300 let data_type = DataType::from(value);
301 let attribute_type = AttributeType::infer(&data_type);
302
303 let mut attribute_data_type = AttributeDataType::new(data_type, attribute_type)
304 .expect("AttributeType was inferred from DataType.");
305
306 match self.0.entry(attribute.clone()) {
307 Entry::Occupied(entry) => {
308 entry.into_mut().merge(&attribute_data_type);
309 }
310 Entry::Vacant(entry) => {
311 if !empty {
312 attribute_data_type.data_type =
313 attribute_data_type.data_type.merge(&DataType::Null);
314 }
315
316 entry.insert(attribute_data_type);
317 }
318 }
319 }
320 }
321
322 #[must_use]
323 pub fn infer(attributes: impl IntoIterator<Item = impl Borrow<AttributeMap>>) -> Self {
324 let mut schema = Self::default();
325
326 let mut empty = true;
327
328 for attributes in attributes {
329 schema.update(attributes.borrow(), empty);
330
331 empty = false;
332 }
333
334 schema
335 }
336}
337
338#[derive(Debug, Clone, PartialEq, Eq, Default)]
339#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
340pub struct GroupSchema {
341 nodes: AttributeSchema,
342 edges: AttributeSchema,
343}
344
345impl GroupSchema {
346 #[must_use]
347 pub const fn new(nodes: AttributeSchema, edges: AttributeSchema) -> Self {
348 Self { nodes, edges }
349 }
350
351 #[must_use]
352 pub fn nodes(&self) -> &AttributeSchemaMapping {
353 &self.nodes
354 }
355
356 #[must_use]
357 pub fn edges(&self) -> &AttributeSchemaMapping {
358 &self.edges
359 }
360
361 pub fn validate_node(
362 &self,
363 index: &NodeIndex,
364 attributes: &AttributeMap,
365 ) -> Result<(), SchemaError> {
366 self.nodes
367 .validate(attributes, &AttributeSchemaKind::Node(index))
368 }
369
370 pub fn validate_edge(
371 &self,
372 index: &EdgeIndex,
373 attributes: &AttributeMap,
374 ) -> Result<(), SchemaError> {
375 self.edges
376 .validate(attributes, &AttributeSchemaKind::Edge(index))
377 }
378
379 #[must_use]
380 pub fn infer(
381 nodes: impl IntoIterator<Item = impl Borrow<AttributeMap>>,
382 edges: impl IntoIterator<Item = impl Borrow<AttributeMap>>,
383 ) -> Self {
384 Self {
385 nodes: AttributeSchema::infer(nodes),
386 edges: AttributeSchema::infer(edges),
387 }
388 }
389
390 pub(crate) fn update_node(&mut self, attributes: &AttributeMap, empty: bool) {
391 self.nodes.update(attributes, empty);
392 }
393
394 pub(crate) fn update_edge(&mut self, attributes: &AttributeMap, empty: bool) {
395 self.edges.update(attributes, empty);
396 }
397}
398
399#[derive(Debug, Clone, PartialEq, Eq, Default)]
400#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
401pub enum SchemaType {
402 #[default]
403 Inferred,
404 Provided,
405}
406
407#[derive(Debug, Clone, PartialEq, Eq, Default)]
408#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
409pub struct Schema {
410 groups: HashMap<Group, GroupSchema>,
411 ungrouped: GroupSchema,
412 r#type: SchemaType,
413}
414
415impl Schema {
416 #[must_use]
417 pub const fn new_inferred(groups: HashMap<Group, GroupSchema>, ungrouped: GroupSchema) -> Self {
418 Self {
419 groups,
420 ungrouped,
421 r#type: SchemaType::Inferred,
422 }
423 }
424
425 #[must_use]
426 pub const fn new_provided(groups: HashMap<Group, GroupSchema>, ungrouped: GroupSchema) -> Self {
427 Self {
428 groups,
429 ungrouped,
430 r#type: SchemaType::Provided,
431 }
432 }
433
434 #[must_use]
435 pub fn infer(graphrecord: &GraphRecord) -> Self {
436 let mut group_mapping: GrHashMap<_, _> = graphrecord
437 .groups()
438 .map(|group| (group, (Vec::new(), Vec::new())))
439 .collect();
440
441 let mut ungrouped = (Vec::new(), Vec::new());
442
443 for node_index in graphrecord.node_indices() {
444 #[expect(clippy::missing_panics_doc, reason = "infallible")]
445 let mut groups_of_node = graphrecord
446 .groups_of_node(node_index)
447 .expect("Node must exist")
448 .peekable();
449
450 if groups_of_node.peek().is_none() {
451 ungrouped.0.push(node_index);
452 continue;
453 }
454
455 for group in groups_of_node {
456 #[expect(clippy::missing_panics_doc, reason = "infallible")]
457 let group_nodes = &mut group_mapping.get_mut(&group).expect("Group must exist").0;
458
459 group_nodes.push(node_index);
460 }
461 }
462
463 for edge_index in graphrecord.edge_indices() {
464 #[expect(clippy::missing_panics_doc, reason = "infallible")]
465 let mut groups_of_edge = graphrecord
466 .groups_of_edge(edge_index)
467 .expect("Edge must exist")
468 .peekable();
469
470 if groups_of_edge.peek().is_none() {
471 ungrouped.1.push(edge_index);
472 continue;
473 }
474
475 for group in groups_of_edge {
476 #[expect(clippy::missing_panics_doc, reason = "infallible")]
477 let group_edges = &mut group_mapping.get_mut(&group).expect("Group must exist").1;
478
479 group_edges.push(edge_index);
480 }
481 }
482
483 #[expect(clippy::missing_panics_doc, reason = "infallible")]
484 let group_schemas =
485 group_mapping
486 .into_iter()
487 .map(|(group, (nodes_in_group, edges_in_group))| {
488 let schema = GroupSchema::infer(
489 nodes_in_group.into_iter().map(|node| {
490 graphrecord.node_attributes(node).expect("Node must exist")
491 }),
492 edges_in_group.into_iter().map(|edge| {
493 graphrecord.edge_attributes(edge).expect("Edge must exist")
494 }),
495 );
496
497 (group.clone(), schema)
498 });
499
500 #[expect(clippy::missing_panics_doc, reason = "infallible")]
501 let ungrouped_schema = GroupSchema::infer(
502 ungrouped
503 .0
504 .into_iter()
505 .map(|node| graphrecord.node_attributes(node).expect("Node must exist")),
506 ungrouped
507 .1
508 .into_iter()
509 .map(|edge| graphrecord.edge_attributes(edge).expect("Edge must exist")),
510 );
511
512 Self {
513 groups: group_schemas.collect(),
514 ungrouped: ungrouped_schema,
515 r#type: SchemaType::Inferred,
516 }
517 }
518
519 #[must_use]
520 pub const fn groups(&self) -> &HashMap<Group, GroupSchema> {
521 &self.groups
522 }
523
524 pub fn group(&self, group: &Group) -> Result<&GroupSchema, SchemaError> {
525 self.groups
526 .get(group)
527 .ok_or_else(|| SchemaError::GroupNotInSchema {
528 group: group.clone(),
529 })
530 }
531
532 #[must_use]
533 pub const fn ungrouped(&self) -> &GroupSchema {
534 &self.ungrouped
535 }
536
537 #[must_use]
538 pub const fn schema_type(&self) -> &SchemaType {
539 &self.r#type
540 }
541
542 pub fn validate_node<'a>(
543 &self,
544 index: &'a NodeIndex,
545 attributes: &'a AttributeMap,
546 group: Option<&'a Group>,
547 ) -> Result<(), SchemaError> {
548 match group {
549 Some(group) => {
550 let schema =
551 self.groups
552 .get(group)
553 .ok_or_else(|| SchemaError::GroupNotInSchema {
554 group: group.clone(),
555 })?;
556
557 schema.validate_node(index, attributes)
558 }
559 None => self.ungrouped.validate_node(index, attributes),
560 }
561 }
562
563 pub fn validate_edge<'a>(
564 &self,
565 index: &'a EdgeIndex,
566 attributes: &'a AttributeMap,
567 group: Option<&'a Group>,
568 ) -> Result<(), SchemaError> {
569 match group {
570 Some(group) => {
571 let schema =
572 self.groups
573 .get(group)
574 .ok_or_else(|| SchemaError::GroupNotInSchema {
575 group: group.clone(),
576 })?;
577
578 schema.validate_edge(index, attributes)
579 }
580 None => self.ungrouped.validate_edge(index, attributes),
581 }
582 }
583
584 pub(crate) fn update_node(
585 &mut self,
586 attributes: &AttributeMap,
587 group: Option<&Group>,
588 empty: bool,
589 ) {
590 match group {
591 Some(group) => {
592 self.groups
593 .entry(group.clone())
594 .or_default()
595 .update_node(attributes, empty);
596 }
597 None => self.ungrouped.update_node(attributes, empty),
598 }
599 }
600
601 pub(crate) fn update_edge(
602 &mut self,
603 attributes: &AttributeMap,
604 group: Option<&Group>,
605 empty: bool,
606 ) {
607 match group {
608 Some(group) => {
609 self.groups
610 .entry(group.clone())
611 .or_default()
612 .update_edge(attributes, empty);
613 }
614 None => self.ungrouped.update_edge(attributes, empty),
615 }
616 }
617
618 pub fn set_node_attribute(
619 &mut self,
620 attribute: &GraphRecordAttribute,
621 data_type: DataType,
622 attribute_type: AttributeType,
623 group: Option<&Group>,
624 ) -> Result<(), SchemaError> {
625 let attribute_data_type = AttributeDataType::new(data_type, attribute_type)?;
626
627 match group {
628 Some(group) => {
629 let group_schema = self.groups.entry(group.clone()).or_default();
630 group_schema
631 .nodes
632 .0
633 .insert(attribute.clone(), attribute_data_type);
634 }
635 None => {
636 self.ungrouped
637 .nodes
638 .0
639 .insert(attribute.clone(), attribute_data_type);
640 }
641 }
642
643 Ok(())
644 }
645
646 pub fn set_edge_attribute(
647 &mut self,
648 attribute: &GraphRecordAttribute,
649 data_type: DataType,
650 attribute_type: AttributeType,
651 group: Option<&Group>,
652 ) -> Result<(), SchemaError> {
653 let attribute_data_type = AttributeDataType::new(data_type, attribute_type)?;
654
655 match group {
656 Some(group) => {
657 let group_schema = self.groups.entry(group.clone()).or_default();
658 group_schema
659 .edges
660 .0
661 .insert(attribute.clone(), attribute_data_type);
662 }
663 None => {
664 self.ungrouped
665 .edges
666 .0
667 .insert(attribute.clone(), attribute_data_type);
668 }
669 }
670
671 Ok(())
672 }
673
674 pub fn update_node_attribute(
675 &mut self,
676 attribute: &GraphRecordAttribute,
677 data_type: DataType,
678 attribute_type: AttributeType,
679 group: Option<&Group>,
680 ) -> Result<(), SchemaError> {
681 let attribute_data_type = AttributeDataType::new(data_type, attribute_type)?;
682
683 match group {
684 Some(group) => {
685 let group_schema = self.groups.entry(group.clone()).or_default();
686 group_schema
687 .nodes
688 .0
689 .entry(attribute.clone())
690 .and_modify(|value| value.merge(&attribute_data_type))
691 .or_insert(attribute_data_type);
692 }
693 None => {
694 self.ungrouped
695 .nodes
696 .0
697 .entry(attribute.clone())
698 .and_modify(|value| value.merge(&attribute_data_type))
699 .or_insert(attribute_data_type);
700 }
701 }
702
703 Ok(())
704 }
705
706 pub fn update_edge_attribute(
707 &mut self,
708 attribute: &GraphRecordAttribute,
709 data_type: DataType,
710 attribute_type: AttributeType,
711 group: Option<&Group>,
712 ) -> Result<(), SchemaError> {
713 let attribute_data_type = AttributeDataType::new(data_type, attribute_type)?;
714
715 match group {
716 Some(group) => {
717 let group_schema = self.groups.entry(group.clone()).or_default();
718 group_schema
719 .edges
720 .0
721 .entry(attribute.clone())
722 .and_modify(|value| value.merge(&attribute_data_type))
723 .or_insert(attribute_data_type);
724 }
725 None => {
726 self.ungrouped
727 .edges
728 .0
729 .entry(attribute.clone())
730 .and_modify(|value| value.merge(&attribute_data_type))
731 .or_insert(attribute_data_type);
732 }
733 }
734
735 Ok(())
736 }
737
738 pub fn remove_node_attribute(
739 &mut self,
740 attribute: &GraphRecordAttribute,
741 group: Option<&Group>,
742 ) {
743 match group {
744 Some(group) => {
745 if let Some(group_schema) = self.groups.get_mut(group) {
746 group_schema.nodes.0.remove(attribute);
747 }
748 }
749 None => {
750 self.ungrouped.nodes.0.remove(attribute);
751 }
752 }
753 }
754
755 pub fn remove_edge_attribute(
756 &mut self,
757 attribute: &GraphRecordAttribute,
758 group: Option<&Group>,
759 ) {
760 match group {
761 Some(group) => {
762 if let Some(group_schema) = self.groups.get_mut(group) {
763 group_schema.edges.0.remove(attribute);
764 }
765 }
766 None => {
767 self.ungrouped.edges.0.remove(attribute);
768 }
769 }
770 }
771
772 pub fn add_group(&mut self, group: Group, schema: GroupSchema) -> Result<(), SchemaError> {
773 if self.groups.contains_key(&group) {
774 return Err(SchemaError::GroupAlreadyInSchema { group });
775 }
776
777 self.groups.insert(group, schema);
778
779 Ok(())
780 }
781
782 pub fn remove_group(&mut self, group: &Group) {
783 self.groups.remove(group);
784 }
785
786 pub const fn freeze(&mut self) {
787 self.r#type = SchemaType::Provided;
788 }
789
790 pub const fn unfreeze(&mut self) {
791 self.r#type = SchemaType::Inferred;
792 }
793}
794
795#[cfg(test)]
796mod test {
797 use super::{AttributeDataType, GroupSchema};
798 use crate::{
799 GraphRecord,
800 graphrecord::{
801 AttributeMap, Schema, SchemaType,
802 datatypes::DataType,
803 schema::{AttributeSchema, AttributeSchemaKind, AttributeType},
804 },
805 };
806 use std::collections::HashMap;
807
808 #[test]
809 fn test_attribute_type_infer() {
810 assert_eq!(
811 AttributeType::infer(&DataType::String),
812 AttributeType::Unstructured
813 );
814 assert_eq!(
815 AttributeType::infer(&DataType::Int),
816 AttributeType::Continuous
817 );
818 assert_eq!(
819 AttributeType::infer(&DataType::Float),
820 AttributeType::Continuous
821 );
822 assert_eq!(
823 AttributeType::infer(&DataType::Bool),
824 AttributeType::Categorical
825 );
826 assert_eq!(
827 AttributeType::infer(&DataType::DateTime),
828 AttributeType::Temporal
829 );
830 assert_eq!(
831 AttributeType::infer(&DataType::Duration),
832 AttributeType::Temporal
833 );
834 assert_eq!(
835 AttributeType::infer(&DataType::Null),
836 AttributeType::Unstructured
837 );
838 assert_eq!(
839 AttributeType::infer(&DataType::Any),
840 AttributeType::Unstructured
841 );
842 assert_eq!(
843 AttributeType::infer(&DataType::Union((
844 Box::new(DataType::Int),
845 Box::new(DataType::Float)
846 ))),
847 AttributeType::Continuous
848 );
849 assert_eq!(
850 AttributeType::infer(&DataType::Option(Box::new(DataType::Int))),
851 AttributeType::Continuous
852 );
853 }
854
855 #[test]
856 fn test_attribute_type_merge() {
857 assert_eq!(
858 AttributeType::Categorical.merge(AttributeType::Unstructured),
859 AttributeType::Unstructured
860 );
861 assert_eq!(
862 AttributeType::Unstructured.merge(AttributeType::Categorical),
863 AttributeType::Unstructured
864 );
865
866 assert_eq!(
867 AttributeType::Categorical.merge(AttributeType::Categorical),
868 AttributeType::Categorical
869 );
870 assert_eq!(
871 AttributeType::Categorical.merge(AttributeType::Continuous),
872 AttributeType::Categorical
873 );
874 assert_eq!(
875 AttributeType::Categorical.merge(AttributeType::Temporal),
876 AttributeType::Categorical
877 );
878
879 assert_eq!(
880 AttributeType::Continuous.merge(AttributeType::Categorical),
881 AttributeType::Categorical
882 );
883 assert_eq!(
884 AttributeType::Temporal.merge(AttributeType::Categorical),
885 AttributeType::Categorical
886 );
887
888 assert_eq!(
889 AttributeType::Continuous.merge(AttributeType::Continuous),
890 AttributeType::Continuous
891 );
892
893 assert_eq!(
894 AttributeType::Temporal.merge(AttributeType::Temporal),
895 AttributeType::Temporal
896 );
897
898 assert_eq!(
899 AttributeType::Continuous.merge(AttributeType::Temporal),
900 AttributeType::Unstructured
901 );
902 assert_eq!(
903 AttributeType::Continuous.merge(AttributeType::Unstructured),
904 AttributeType::Unstructured
905 );
906
907 assert_eq!(
908 AttributeType::Temporal.merge(AttributeType::Continuous),
909 AttributeType::Unstructured
910 );
911 assert_eq!(
912 AttributeType::Temporal.merge(AttributeType::Unstructured),
913 AttributeType::Unstructured
914 );
915
916 assert_eq!(
917 AttributeType::Unstructured.merge(AttributeType::Continuous),
918 AttributeType::Unstructured
919 );
920 assert_eq!(
921 AttributeType::Unstructured.merge(AttributeType::Temporal),
922 AttributeType::Unstructured
923 );
924 assert_eq!(
925 AttributeType::Unstructured.merge(AttributeType::Unstructured),
926 AttributeType::Unstructured
927 );
928 }
929
930 #[test]
931 fn test_data_type_merge() {
932 assert_eq!(DataType::Int.merge(&DataType::Int), DataType::Int);
933 assert_eq!(
934 DataType::Int.merge(&DataType::Float),
935 DataType::Union((Box::new(DataType::Int), Box::new(DataType::Float)))
936 );
937 assert_eq!(
938 DataType::Int.merge(&DataType::Null),
939 DataType::Option(Box::new(DataType::Int))
940 );
941 assert_eq!(
942 DataType::Null.merge(&DataType::Int),
943 DataType::Option(Box::new(DataType::Int))
944 );
945 assert_eq!(DataType::Null.merge(&DataType::Null), DataType::Null);
946 assert_eq!(DataType::Int.merge(&DataType::Any), DataType::Any);
947 assert_eq!(DataType::Any.merge(&DataType::Int), DataType::Any);
948 assert_eq!(
949 DataType::Option(Box::new(DataType::Int)).merge(&DataType::String),
950 DataType::Option(Box::new(DataType::Union((
951 Box::new(DataType::Int),
952 Box::new(DataType::String)
953 ))))
954 );
955 assert_eq!(
956 DataType::Int.merge(&DataType::Option(Box::new(DataType::Int))),
957 DataType::Option(Box::new(DataType::Int))
958 );
959 assert_eq!(
960 DataType::Option(Box::new(DataType::Int))
961 .merge(&DataType::Option(Box::new(DataType::String))),
962 DataType::Option(Box::new(DataType::Union((
963 Box::new(DataType::Int),
964 Box::new(DataType::String)
965 ))))
966 );
967 }
968
969 #[test]
970 fn test_attribute_data_type_new() {
971 assert!(AttributeDataType::new(DataType::String, AttributeType::Categorical).is_ok());
972 assert!(AttributeDataType::new(DataType::String, AttributeType::Continuous).is_err());
973 assert!(AttributeDataType::new(DataType::String, AttributeType::Temporal).is_err());
974 assert!(AttributeDataType::new(DataType::String, AttributeType::Unstructured).is_ok());
975
976 assert!(AttributeDataType::new(DataType::Int, AttributeType::Categorical).is_ok());
977 assert!(AttributeDataType::new(DataType::Int, AttributeType::Continuous).is_ok());
978 assert!(AttributeDataType::new(DataType::Int, AttributeType::Temporal).is_err());
979 assert!(AttributeDataType::new(DataType::Int, AttributeType::Unstructured).is_ok());
980
981 assert!(AttributeDataType::new(DataType::Float, AttributeType::Categorical).is_ok());
982 assert!(AttributeDataType::new(DataType::Float, AttributeType::Continuous).is_ok());
983 assert!(AttributeDataType::new(DataType::Float, AttributeType::Temporal).is_err());
984 assert!(AttributeDataType::new(DataType::Float, AttributeType::Unstructured).is_ok());
985
986 assert!(AttributeDataType::new(DataType::Bool, AttributeType::Categorical).is_ok());
987 assert!(AttributeDataType::new(DataType::Bool, AttributeType::Continuous).is_err());
988 assert!(AttributeDataType::new(DataType::Bool, AttributeType::Temporal).is_err());
989 assert!(AttributeDataType::new(DataType::Bool, AttributeType::Unstructured).is_ok());
990
991 assert!(AttributeDataType::new(DataType::DateTime, AttributeType::Categorical).is_ok());
992 assert!(AttributeDataType::new(DataType::DateTime, AttributeType::Continuous).is_err());
993 assert!(AttributeDataType::new(DataType::DateTime, AttributeType::Temporal).is_ok());
994 assert!(AttributeDataType::new(DataType::DateTime, AttributeType::Unstructured).is_ok());
995
996 assert!(AttributeDataType::new(DataType::Duration, AttributeType::Categorical).is_ok());
997 assert!(AttributeDataType::new(DataType::Duration, AttributeType::Continuous).is_err());
998 assert!(AttributeDataType::new(DataType::Duration, AttributeType::Temporal).is_ok());
999 assert!(AttributeDataType::new(DataType::Duration, AttributeType::Unstructured).is_ok());
1000
1001 assert!(AttributeDataType::new(DataType::Null, AttributeType::Categorical).is_ok());
1002 assert!(AttributeDataType::new(DataType::Null, AttributeType::Continuous).is_ok());
1003 assert!(AttributeDataType::new(DataType::Null, AttributeType::Temporal).is_ok());
1004 assert!(AttributeDataType::new(DataType::Null, AttributeType::Unstructured).is_ok());
1005
1006 assert!(AttributeDataType::new(DataType::Any, AttributeType::Categorical).is_ok());
1007 assert!(AttributeDataType::new(DataType::Any, AttributeType::Continuous).is_err());
1008 assert!(AttributeDataType::new(DataType::Any, AttributeType::Temporal).is_err());
1009 assert!(AttributeDataType::new(DataType::Any, AttributeType::Unstructured).is_ok());
1010
1011 assert!(
1012 AttributeDataType::new(
1013 DataType::Option(Box::new(DataType::Int)),
1014 AttributeType::Categorical
1015 )
1016 .is_ok()
1017 );
1018 assert!(
1019 AttributeDataType::new(
1020 DataType::Option(Box::new(DataType::Int)),
1021 AttributeType::Continuous
1022 )
1023 .is_ok()
1024 );
1025 assert!(
1026 AttributeDataType::new(
1027 DataType::Option(Box::new(DataType::Int)),
1028 AttributeType::Temporal
1029 )
1030 .is_err()
1031 );
1032 assert!(
1033 AttributeDataType::new(
1034 DataType::Option(Box::new(DataType::Int)),
1035 AttributeType::Unstructured
1036 )
1037 .is_ok()
1038 );
1039
1040 assert!(
1041 AttributeDataType::new(
1042 DataType::Union((Box::new(DataType::Int), Box::new(DataType::Float))),
1043 AttributeType::Categorical
1044 )
1045 .is_ok()
1046 );
1047 assert!(
1048 AttributeDataType::new(
1049 DataType::Union((Box::new(DataType::Int), Box::new(DataType::Float))),
1050 AttributeType::Continuous
1051 )
1052 .is_ok()
1053 );
1054 assert!(
1055 AttributeDataType::new(
1056 DataType::Union((Box::new(DataType::Int), Box::new(DataType::Float))),
1057 AttributeType::Temporal
1058 )
1059 .is_err()
1060 );
1061 assert!(
1062 AttributeDataType::new(
1063 DataType::Union((Box::new(DataType::Int), Box::new(DataType::Float))),
1064 AttributeType::Unstructured
1065 )
1066 .is_ok()
1067 );
1068 }
1069
1070 #[test]
1071 fn test_attribute_data_type_data_type() {
1072 let attribute_data_type = AttributeDataType::new(DataType::Int, AttributeType::Categorical)
1073 .expect("AttributeType was inferred from DataType.");
1074
1075 assert_eq!(attribute_data_type.data_type(), &DataType::Int);
1076 }
1077
1078 #[test]
1079 fn test_attribute_data_type_attribute_type() {
1080 let attribute_data_type = AttributeDataType::new(DataType::Int, AttributeType::Categorical)
1081 .expect("AttributeType was inferred from DataType.");
1082
1083 assert_eq!(
1084 attribute_data_type.attribute_type(),
1085 &AttributeType::Categorical
1086 );
1087 }
1088
1089 #[test]
1090 fn test_attribute_data_type_merge() {
1091 let mut attribute_data_type =
1092 AttributeDataType::new(DataType::Int, AttributeType::Categorical)
1093 .expect("AttributeType was inferred from DataType.");
1094
1095 attribute_data_type.merge(
1096 &AttributeDataType::new(DataType::Float, AttributeType::Continuous)
1097 .expect("AttributeType was inferred from DataType."),
1098 );
1099
1100 assert_eq!(
1101 attribute_data_type.data_type(),
1102 &DataType::Union((Box::new(DataType::Int), Box::new(DataType::Float)))
1103 );
1104 assert_eq!(
1105 attribute_data_type.attribute_type(),
1106 &AttributeType::Categorical
1107 );
1108 }
1109
1110 #[test]
1111 fn test_attribute_data_type_from_data_type() {
1112 let attribute_data_type: AttributeDataType = DataType::Int.into();
1113
1114 assert_eq!(attribute_data_type.data_type(), &DataType::Int);
1115 assert_eq!(
1116 attribute_data_type.attribute_type(),
1117 &AttributeType::Continuous
1118 );
1119 }
1120
1121 #[test]
1122 fn test_attribute_data_type_from_tuple() {
1123 let attribute_data_type: AttributeDataType =
1124 (DataType::Int, AttributeType::Categorical).into();
1125
1126 assert_eq!(attribute_data_type.data_type(), &DataType::Int);
1127 assert_eq!(
1128 attribute_data_type.attribute_type(),
1129 &AttributeType::Categorical
1130 );
1131 }
1132
1133 #[test]
1134 fn test_attribute_schema_deref() {
1135 let schema = AttributeSchema::new(
1136 vec![
1137 (
1138 "key1".into(),
1139 AttributeDataType::new(DataType::Int, AttributeType::Categorical)
1140 .expect("AttributeType was inferred from DataType."),
1141 ),
1142 (
1143 "key2".into(),
1144 AttributeDataType::new(DataType::Float, AttributeType::Continuous)
1145 .expect("AttributeType was inferred from DataType."),
1146 ),
1147 ]
1148 .into_iter()
1149 .collect(),
1150 );
1151
1152 assert_eq!(
1153 schema.get(&"key1".into()).unwrap().data_type(),
1154 &DataType::Int
1155 );
1156 assert_eq!(
1157 schema.get(&"key2".into()).unwrap().data_type(),
1158 &DataType::Float
1159 );
1160 }
1161
1162 #[test]
1163 fn test_attribute_schema_validate() {
1164 let attribute_schema = AttributeSchema::new(
1165 vec![
1166 (
1167 "key1".into(),
1168 AttributeDataType::new(DataType::Int, AttributeType::Categorical)
1169 .expect("AttributeType was inferred from DataType."),
1170 ),
1171 (
1172 "key2".into(),
1173 AttributeDataType::new(DataType::Float, AttributeType::Continuous)
1174 .expect("AttributeType was inferred from DataType."),
1175 ),
1176 ]
1177 .into_iter()
1178 .collect(),
1179 );
1180
1181 let attributes: AttributeMap = vec![("key1".into(), 0.into()), ("key2".into(), 0.0.into())]
1182 .into_iter()
1183 .collect();
1184
1185 assert!(
1186 attribute_schema
1187 .validate(&attributes, &AttributeSchemaKind::Node(&0.into()))
1188 .is_ok()
1189 );
1190
1191 let attributes: AttributeMap = vec![("key1".into(), 0.0.into()), ("key2".into(), 0.into())]
1192 .into_iter()
1193 .collect();
1194
1195 assert!(
1196 attribute_schema
1197 .validate(&attributes, &AttributeSchemaKind::Node(&0.into()))
1198 .is_err_and(|error| {
1199 matches!(
1200 error,
1201 crate::errors::SchemaError::NodeAttributeDataTypeMismatch { .. }
1202 )
1203 })
1204 );
1205
1206 let attributes: AttributeMap = vec![
1207 ("key1".into(), 0.into()),
1208 ("key2".into(), 0.0.into()),
1209 ("key3".into(), 0.0.into()),
1210 ]
1211 .into_iter()
1212 .collect();
1213
1214 assert!(
1215 attribute_schema
1216 .validate(&attributes, &AttributeSchemaKind::Node(&0.into()))
1217 .is_err_and(|error| {
1218 matches!(
1219 error,
1220 crate::errors::SchemaError::NodeAttributesNotInSchema { .. }
1221 )
1222 })
1223 );
1224 }
1225
1226 #[test]
1227 fn test_attribute_schema_update() {
1228 let mut schema = AttributeSchema::default();
1229 let attributes: AttributeMap =
1230 vec![("key1".into(), 0.into()), ("key2".into(), "test".into())]
1231 .into_iter()
1232 .collect();
1233
1234 schema.update(&attributes, true);
1235
1236 assert_eq!(schema.0.len(), 2);
1237 assert_eq!(
1238 schema.0.get(&"key1".into()).unwrap().data_type(),
1239 &DataType::Int
1240 );
1241 assert_eq!(
1242 schema.0.get(&"key2".into()).unwrap().data_type(),
1243 &DataType::String
1244 );
1245
1246 let new_attributes: AttributeMap =
1247 vec![("key1".into(), 0.5.into()), ("key3".into(), true.into())]
1248 .into_iter()
1249 .collect();
1250
1251 schema.update(&new_attributes, false);
1252
1253 assert_eq!(schema.0.len(), 3);
1254 assert_eq!(
1255 schema.0.get(&"key1".into()).unwrap().data_type(),
1256 &DataType::Union((Box::new(DataType::Int), Box::new(DataType::Float)))
1257 );
1258 assert_eq!(
1259 schema.0.get(&"key2".into()).unwrap().data_type(),
1260 &DataType::Option(Box::new(DataType::String))
1261 );
1262 assert_eq!(
1263 schema.0.get(&"key3".into()).unwrap().data_type(),
1264 &DataType::Option(Box::new(DataType::Bool))
1265 );
1266 }
1267
1268 #[test]
1269 fn test_attribute_schema_infer() {
1270 let attributes1: AttributeMap =
1271 vec![("key1".into(), 0.into()), ("key2".into(), "test".into())]
1272 .into_iter()
1273 .collect();
1274
1275 let attributes2: AttributeMap =
1276 vec![("key1".into(), 1.into()), ("key3".into(), true.into())]
1277 .into_iter()
1278 .collect();
1279
1280 let schema = AttributeSchema::infer(vec![&attributes1, &attributes2]);
1281
1282 assert_eq!(schema.0.len(), 3);
1283 assert_eq!(
1284 schema.0.get(&"key1".into()).unwrap().data_type(),
1285 &DataType::Int
1286 );
1287 assert_eq!(
1288 schema.0.get(&"key2".into()).unwrap().data_type(),
1289 &DataType::Option(Box::new(DataType::String))
1290 );
1291 assert_eq!(
1292 schema.0.get(&"key3".into()).unwrap().data_type(),
1293 &DataType::Option(Box::new(DataType::Bool))
1294 );
1295 }
1296
1297 #[test]
1298 fn test_group_schema_nodes() {
1299 let nodes = AttributeSchema::new(
1300 vec![
1301 (
1302 "key1".into(),
1303 AttributeDataType::new(DataType::Int, AttributeType::Categorical)
1304 .expect("AttributeType was inferred from DataType."),
1305 ),
1306 (
1307 "key2".into(),
1308 AttributeDataType::new(DataType::Float, AttributeType::Continuous)
1309 .expect("AttributeType was inferred from DataType."),
1310 ),
1311 ]
1312 .into_iter()
1313 .collect(),
1314 );
1315
1316 let group_schema = GroupSchema::new(nodes.clone(), AttributeSchema::default());
1317
1318 assert_eq!(group_schema.nodes(), &nodes.0);
1319 }
1320
1321 #[test]
1322 fn test_group_schema_edges() {
1323 let edges = AttributeSchema::new(
1324 vec![
1325 (
1326 "key1".into(),
1327 AttributeDataType::new(DataType::Int, AttributeType::Categorical)
1328 .expect("AttributeType was inferred from DataType."),
1329 ),
1330 (
1331 "key2".into(),
1332 AttributeDataType::new(DataType::Float, AttributeType::Continuous)
1333 .expect("AttributeType was inferred from DataType."),
1334 ),
1335 ]
1336 .into_iter()
1337 .collect(),
1338 );
1339
1340 let group_schema = GroupSchema::new(AttributeSchema::default(), edges.clone());
1341
1342 assert_eq!(group_schema.edges(), &edges.0);
1343 }
1344
1345 #[test]
1346 fn test_group_schema_validate_node() {
1347 let nodes = AttributeSchema::new(
1348 vec![
1349 (
1350 "key1".into(),
1351 AttributeDataType::new(DataType::Int, AttributeType::Categorical)
1352 .expect("AttributeType was inferred from DataType."),
1353 ),
1354 (
1355 "key2".into(),
1356 AttributeDataType::new(DataType::Float, AttributeType::Continuous)
1357 .expect("AttributeType was inferred from DataType."),
1358 ),
1359 ]
1360 .into_iter()
1361 .collect(),
1362 );
1363
1364 let group_schema = GroupSchema::new(nodes, AttributeSchema::default());
1365
1366 let attributes: AttributeMap = vec![("key1".into(), 0.into()), ("key2".into(), 0.0.into())]
1367 .into_iter()
1368 .collect();
1369
1370 assert!(group_schema.validate_node(&0.into(), &attributes).is_ok());
1371
1372 let attributes: AttributeMap = vec![("key1".into(), 0.0.into()), ("key2".into(), 0.into())]
1373 .into_iter()
1374 .collect();
1375
1376 assert!(
1377 group_schema
1378 .validate_node(&0.into(), &attributes)
1379 .is_err_and(|error| {
1380 matches!(
1381 error,
1382 crate::errors::SchemaError::NodeAttributeDataTypeMismatch { .. }
1383 )
1384 })
1385 );
1386 }
1387
1388 #[test]
1389 fn test_group_schema_validate_edge() {
1390 let edges = AttributeSchema::new(
1391 vec![
1392 (
1393 "key1".into(),
1394 AttributeDataType::new(DataType::Int, AttributeType::Categorical)
1395 .expect("AttributeType was inferred from DataType."),
1396 ),
1397 (
1398 "key2".into(),
1399 AttributeDataType::new(DataType::Float, AttributeType::Continuous)
1400 .expect("AttributeType was inferred from DataType."),
1401 ),
1402 ]
1403 .into_iter()
1404 .collect(),
1405 );
1406
1407 let group_schema = GroupSchema::new(AttributeSchema::default(), edges);
1408
1409 let attributes: AttributeMap = vec![("key1".into(), 0.into()), ("key2".into(), 0.0.into())]
1410 .into_iter()
1411 .collect();
1412
1413 assert!(group_schema.validate_edge(&0, &attributes).is_ok());
1414
1415 let attributes: AttributeMap = vec![("key1".into(), 0.0.into()), ("key2".into(), 0.into())]
1416 .into_iter()
1417 .collect();
1418
1419 assert!(
1420 group_schema
1421 .validate_edge(&0, &attributes)
1422 .is_err_and(|error| {
1423 matches!(
1424 error,
1425 crate::errors::SchemaError::EdgeAttributeDataTypeMismatch { .. }
1426 )
1427 })
1428 );
1429 }
1430
1431 #[test]
1432 fn test_group_schema_infer() {
1433 let node_attributes1: AttributeMap =
1434 vec![("key1".into(), 0.into()), ("key2".into(), "test".into())]
1435 .into_iter()
1436 .collect();
1437
1438 let node_attributes2: AttributeMap =
1439 vec![("key1".into(), 1.into()), ("key3".into(), true.into())]
1440 .into_iter()
1441 .collect();
1442
1443 let edge_attributes: AttributeMap =
1444 vec![("key4".into(), 0.5.into()), ("key5".into(), "edge".into())]
1445 .into_iter()
1446 .collect();
1447
1448 let group_schema = GroupSchema::infer(
1449 vec![&node_attributes1, &node_attributes2],
1450 vec![&edge_attributes],
1451 );
1452
1453 assert_eq!(group_schema.nodes().len(), 3);
1454 assert_eq!(group_schema.edges().len(), 2);
1455
1456 assert_eq!(
1457 group_schema
1458 .nodes()
1459 .get(&"key1".into())
1460 .unwrap()
1461 .data_type(),
1462 &DataType::Int
1463 );
1464 assert_eq!(
1465 group_schema
1466 .nodes()
1467 .get(&"key2".into())
1468 .unwrap()
1469 .data_type(),
1470 &DataType::Option(Box::new(DataType::String))
1471 );
1472 assert_eq!(
1473 group_schema
1474 .nodes()
1475 .get(&"key3".into())
1476 .unwrap()
1477 .data_type(),
1478 &DataType::Option(Box::new(DataType::Bool))
1479 );
1480
1481 assert_eq!(
1482 group_schema
1483 .edges()
1484 .get(&"key4".into())
1485 .unwrap()
1486 .data_type(),
1487 &DataType::Float
1488 );
1489 assert_eq!(
1490 group_schema
1491 .edges()
1492 .get(&"key5".into())
1493 .unwrap()
1494 .data_type(),
1495 &DataType::String
1496 );
1497 }
1498
1499 #[test]
1500 fn test_group_schema_update_node() {
1501 let mut group_schema = GroupSchema::default();
1502 let attributes =
1503 AttributeMap::from([("key1".into(), 0.into()), ("key2".into(), 0.0.into())]);
1504
1505 group_schema.update_node(&attributes, true);
1506
1507 assert_eq!(group_schema.nodes().len(), 2);
1508 assert_eq!(
1509 group_schema
1510 .nodes()
1511 .get(&"key1".into())
1512 .unwrap()
1513 .data_type(),
1514 &DataType::Int
1515 );
1516 assert_eq!(
1517 group_schema
1518 .nodes()
1519 .get(&"key2".into())
1520 .unwrap()
1521 .data_type(),
1522 &DataType::Float
1523 );
1524 }
1525
1526 #[test]
1527 fn test_group_schema_update_edge() {
1528 let mut group_schema = GroupSchema::default();
1529 let attributes =
1530 AttributeMap::from([("key3".into(), true.into()), ("key4".into(), "test".into())]);
1531
1532 group_schema.update_edge(&attributes, true);
1533
1534 assert_eq!(group_schema.edges().len(), 2);
1535 assert_eq!(
1536 group_schema
1537 .edges()
1538 .get(&"key3".into())
1539 .unwrap()
1540 .data_type(),
1541 &DataType::Bool
1542 );
1543 assert_eq!(
1544 group_schema
1545 .edges()
1546 .get(&"key4".into())
1547 .unwrap()
1548 .data_type(),
1549 &DataType::String
1550 );
1551 }
1552
1553 #[test]
1554 fn test_schema_infer() {
1555 let mut graphrecord = GraphRecord::new();
1556 graphrecord
1557 .add_node(0.into(), AttributeMap::from([("key1".into(), 0.into())]))
1558 .unwrap();
1559 graphrecord
1560 .add_node(1.into(), AttributeMap::from([("key2".into(), 0.0.into())]))
1561 .unwrap();
1562 graphrecord
1563 .add_edge(
1564 0.into(),
1565 1.into(),
1566 AttributeMap::from([("key3".into(), true.into())]),
1567 )
1568 .unwrap();
1569
1570 let schema = Schema::infer(&graphrecord);
1571
1572 assert_eq!(schema.ungrouped().nodes().len(), 2);
1573 assert_eq!(schema.ungrouped().edges().len(), 1);
1574
1575 graphrecord
1576 .add_group("test".into(), Some(vec![0.into(), 1.into()]), Some(vec![0]))
1577 .unwrap();
1578
1579 let schema = Schema::infer(&graphrecord);
1580
1581 assert_eq!(schema.groups().len(), 1);
1582 assert_eq!(schema.group(&"test".into()).unwrap().nodes().len(), 2);
1583 assert_eq!(schema.group(&"test".into()).unwrap().edges().len(), 1);
1584 }
1585
1586 #[test]
1587 fn test_schema_groups() {
1588 let schema = Schema::new_inferred(
1589 vec![("group1".into(), GroupSchema::default())]
1590 .into_iter()
1591 .collect(),
1592 GroupSchema::default(),
1593 );
1594 assert_eq!(schema.groups().len(), 1);
1595 assert!(schema.groups().contains_key(&"group1".into()));
1596 }
1597
1598 #[test]
1599 fn test_schema_group() {
1600 let schema = Schema::new_inferred(
1601 vec![("group1".into(), GroupSchema::default())]
1602 .into_iter()
1603 .collect(),
1604 GroupSchema::default(),
1605 );
1606 assert!(schema.group(&"group1".into()).is_ok());
1607 assert!(schema.group(&"non_existent".into()).is_err());
1608 }
1609
1610 #[test]
1611 fn test_schema_default() {
1612 let default_schema = GroupSchema::default();
1613 let schema = Schema::new_inferred(HashMap::new(), default_schema.clone());
1614 assert_eq!(schema.ungrouped(), &default_schema);
1615 }
1616
1617 #[test]
1618 fn test_schema_schema_type() {
1619 let schema = Schema::new_inferred(HashMap::new(), GroupSchema::default());
1620 assert_eq!(schema.schema_type(), &SchemaType::Inferred);
1621 }
1622
1623 #[test]
1624 fn test_schema_validate_node() {
1625 let mut schema = Schema::new_inferred(
1626 HashMap::new(),
1627 GroupSchema::new(AttributeSchema::default(), AttributeSchema::default()),
1628 );
1629 schema
1630 .set_node_attribute(
1631 &"key1".into(),
1632 DataType::Int,
1633 AttributeType::Continuous,
1634 None,
1635 )
1636 .unwrap();
1637
1638 let attributes = AttributeMap::from([("key1".into(), 0.into())]);
1639 assert!(schema.validate_node(&0.into(), &attributes, None).is_ok());
1640
1641 let invalid_attributes = AttributeMap::from([("key1".into(), "invalid".into())]);
1642 assert!(
1643 schema
1644 .validate_node(&0.into(), &invalid_attributes, None)
1645 .is_err()
1646 );
1647 }
1648
1649 #[test]
1650 fn test_schema_validate_edge() {
1651 let mut schema = Schema::new_inferred(
1652 HashMap::new(),
1653 GroupSchema::new(AttributeSchema::default(), AttributeSchema::default()),
1654 );
1655 schema
1656 .set_edge_attribute(
1657 &"key1".into(),
1658 DataType::Bool,
1659 AttributeType::Categorical,
1660 None,
1661 )
1662 .unwrap();
1663
1664 let attributes = AttributeMap::from([("key1".into(), true.into())]);
1665 assert!(schema.validate_edge(&0, &attributes, None).is_ok());
1666
1667 let invalid_attributes = AttributeMap::from([("key1".into(), 0.into())]);
1668 assert!(schema.validate_edge(&0, &invalid_attributes, None).is_err());
1669 }
1670
1671 #[test]
1672 fn test_schema_update_node() {
1673 let mut schema = Schema::new_inferred(
1674 HashMap::new(),
1675 GroupSchema::new(AttributeSchema::default(), AttributeSchema::default()),
1676 );
1677 let attributes =
1678 AttributeMap::from([("key1".into(), 0.into()), ("key2".into(), 0.0.into())]);
1679
1680 schema.update_node(&attributes, None, true);
1681
1682 assert_eq!(schema.ungrouped().nodes().len(), 2);
1683 assert_eq!(
1684 schema
1685 .ungrouped()
1686 .nodes()
1687 .get(&"key1".into())
1688 .unwrap()
1689 .data_type(),
1690 &DataType::Int
1691 );
1692 assert_eq!(
1693 schema
1694 .ungrouped()
1695 .nodes()
1696 .get(&"key2".into())
1697 .unwrap()
1698 .data_type(),
1699 &DataType::Float
1700 );
1701 }
1702
1703 #[test]
1704 fn test_schema_update_edge() {
1705 let mut schema = Schema::new_inferred(
1706 HashMap::new(),
1707 GroupSchema::new(AttributeSchema::default(), AttributeSchema::default()),
1708 );
1709 let attributes =
1710 AttributeMap::from([("key3".into(), true.into()), ("key4".into(), "test".into())]);
1711
1712 schema.update_edge(&attributes, None, true);
1713
1714 assert_eq!(schema.ungrouped().edges().len(), 2);
1715 assert_eq!(
1716 schema
1717 .ungrouped()
1718 .edges()
1719 .get(&"key3".into())
1720 .unwrap()
1721 .data_type(),
1722 &DataType::Bool
1723 );
1724 assert_eq!(
1725 schema
1726 .ungrouped()
1727 .edges()
1728 .get(&"key4".into())
1729 .unwrap()
1730 .data_type(),
1731 &DataType::String
1732 );
1733 }
1734
1735 #[test]
1736 fn test_schema_set_node_attribute() {
1737 let mut schema = Schema::new_inferred(HashMap::new(), GroupSchema::default());
1738 assert!(
1739 schema
1740 .set_node_attribute(
1741 &"key1".into(),
1742 DataType::Int,
1743 AttributeType::Continuous,
1744 None
1745 )
1746 .is_ok()
1747 );
1748 assert_eq!(
1749 schema
1750 .ungrouped()
1751 .nodes()
1752 .get(&"key1".into())
1753 .unwrap()
1754 .data_type(),
1755 &DataType::Int
1756 );
1757 assert!(
1758 schema
1759 .set_node_attribute(
1760 &"key1".into(),
1761 DataType::Float,
1762 AttributeType::Continuous,
1763 None
1764 )
1765 .is_ok()
1766 );
1767 assert_eq!(
1768 schema
1769 .ungrouped()
1770 .nodes()
1771 .get(&"key1".into())
1772 .unwrap()
1773 .data_type(),
1774 &DataType::Float
1775 );
1776
1777 assert!(
1778 schema
1779 .set_node_attribute(
1780 &"key1".into(),
1781 DataType::Float,
1782 AttributeType::Continuous,
1783 Some(&"group1".into())
1784 )
1785 .is_ok()
1786 );
1787 assert_eq!(
1788 schema
1789 .group(&"group1".into())
1790 .unwrap()
1791 .nodes()
1792 .get(&"key1".into())
1793 .unwrap()
1794 .data_type(),
1795 &DataType::Float
1796 );
1797 }
1798
1799 #[test]
1800 fn test_schema_set_edge_attribute() {
1801 let mut schema = Schema::new_inferred(HashMap::new(), GroupSchema::default());
1802 assert!(
1803 schema
1804 .set_edge_attribute(
1805 &"key1".into(),
1806 DataType::Bool,
1807 AttributeType::Categorical,
1808 None
1809 )
1810 .is_ok()
1811 );
1812 assert_eq!(
1813 schema
1814 .ungrouped()
1815 .edges()
1816 .get(&"key1".into())
1817 .unwrap()
1818 .data_type(),
1819 &DataType::Bool
1820 );
1821 assert!(
1822 schema
1823 .set_edge_attribute(
1824 &"key1".into(),
1825 DataType::Float,
1826 AttributeType::Continuous,
1827 None
1828 )
1829 .is_ok()
1830 );
1831 assert_eq!(
1832 schema
1833 .ungrouped()
1834 .edges()
1835 .get(&"key1".into())
1836 .unwrap()
1837 .data_type(),
1838 &DataType::Float
1839 );
1840
1841 assert!(
1842 schema
1843 .set_edge_attribute(
1844 &"key1".into(),
1845 DataType::Float,
1846 AttributeType::Continuous,
1847 Some(&"group1".into())
1848 )
1849 .is_ok()
1850 );
1851 assert_eq!(
1852 schema
1853 .group(&"group1".into())
1854 .unwrap()
1855 .edges()
1856 .get(&"key1".into())
1857 .unwrap()
1858 .data_type(),
1859 &DataType::Float
1860 );
1861 }
1862
1863 #[test]
1864 fn test_schema_update_node_attribute() {
1865 let mut schema = Schema::new_inferred(HashMap::new(), GroupSchema::default());
1866 schema
1867 .set_node_attribute(
1868 &"key1".into(),
1869 DataType::Int,
1870 AttributeType::Continuous,
1871 None,
1872 )
1873 .unwrap();
1874 assert!(
1875 schema
1876 .update_node_attribute(
1877 &"key1".into(),
1878 DataType::Float,
1879 AttributeType::Continuous,
1880 None
1881 )
1882 .is_ok()
1883 );
1884 assert_eq!(
1885 schema
1886 .ungrouped()
1887 .nodes()
1888 .get(&"key1".into())
1889 .unwrap()
1890 .data_type(),
1891 &DataType::Union((Box::new(DataType::Int), Box::new(DataType::Float)))
1892 );
1893
1894 schema
1895 .set_node_attribute(
1896 &"key1".into(),
1897 DataType::Int,
1898 AttributeType::Continuous,
1899 Some(&"group1".into()),
1900 )
1901 .unwrap();
1902 assert!(
1903 schema
1904 .update_node_attribute(
1905 &"key1".into(),
1906 DataType::Float,
1907 AttributeType::Continuous,
1908 Some(&"group1".into())
1909 )
1910 .is_ok()
1911 );
1912 assert_eq!(
1913 schema
1914 .group(&"group1".into())
1915 .unwrap()
1916 .nodes()
1917 .get(&"key1".into())
1918 .unwrap()
1919 .data_type(),
1920 &DataType::Union((Box::new(DataType::Int), Box::new(DataType::Float)))
1921 );
1922 }
1923
1924 #[test]
1925 fn test_schema_update_edge_attribute() {
1926 let mut schema = Schema::new_inferred(HashMap::new(), GroupSchema::default());
1927 schema
1928 .set_edge_attribute(
1929 &"key1".into(),
1930 DataType::Bool,
1931 AttributeType::Categorical,
1932 None,
1933 )
1934 .unwrap();
1935 assert!(
1936 schema
1937 .update_edge_attribute(
1938 &"key1".into(),
1939 DataType::String,
1940 AttributeType::Unstructured,
1941 None
1942 )
1943 .is_ok()
1944 );
1945 assert_eq!(
1946 schema
1947 .ungrouped()
1948 .edges()
1949 .get(&"key1".into())
1950 .unwrap()
1951 .data_type(),
1952 &DataType::Union((Box::new(DataType::Bool), Box::new(DataType::String)))
1953 );
1954
1955 schema
1956 .set_edge_attribute(
1957 &"key1".into(),
1958 DataType::Bool,
1959 AttributeType::Categorical,
1960 Some(&"group1".into()),
1961 )
1962 .unwrap();
1963 assert!(
1964 schema
1965 .update_edge_attribute(
1966 &"key1".into(),
1967 DataType::String,
1968 AttributeType::Unstructured,
1969 Some(&"group1".into())
1970 )
1971 .is_ok()
1972 );
1973 assert_eq!(
1974 schema
1975 .group(&"group1".into())
1976 .unwrap()
1977 .edges()
1978 .get(&"key1".into())
1979 .unwrap()
1980 .data_type(),
1981 &DataType::Union((Box::new(DataType::Bool), Box::new(DataType::String)))
1982 );
1983 }
1984
1985 #[test]
1986 fn test_schema_remove_node_attribute() {
1987 let mut schema = Schema::new_inferred(HashMap::new(), GroupSchema::default());
1988 schema
1989 .set_node_attribute(
1990 &"key1".into(),
1991 DataType::Int,
1992 AttributeType::Continuous,
1993 None,
1994 )
1995 .unwrap();
1996 schema.remove_node_attribute(&"key1".into(), None);
1997 assert!(!schema.ungrouped().nodes().contains_key(&"key1".into()));
1998
1999 schema
2000 .set_node_attribute(
2001 &"key1".into(),
2002 DataType::Int,
2003 AttributeType::Continuous,
2004 Some(&"group1".into()),
2005 )
2006 .unwrap();
2007 schema.remove_node_attribute(&"key1".into(), Some(&"group1".into()));
2008 assert!(
2009 !schema
2010 .group(&"group1".into())
2011 .unwrap()
2012 .nodes()
2013 .contains_key(&"key1".into())
2014 );
2015 }
2016
2017 #[test]
2018 fn test_schema_remove_edge_attribute() {
2019 let mut schema = Schema::new_inferred(HashMap::new(), GroupSchema::default());
2020 schema
2021 .set_edge_attribute(
2022 &"key1".into(),
2023 DataType::Bool,
2024 AttributeType::Categorical,
2025 None,
2026 )
2027 .unwrap();
2028 schema.remove_edge_attribute(&"key1".into(), None);
2029 assert!(!schema.ungrouped().edges().contains_key(&"key1".into()));
2030
2031 schema
2032 .set_edge_attribute(
2033 &"key1".into(),
2034 DataType::Bool,
2035 AttributeType::Categorical,
2036 Some(&"group1".into()),
2037 )
2038 .unwrap();
2039 schema.remove_edge_attribute(&"key1".into(), Some(&"group1".into()));
2040 assert!(
2041 !schema
2042 .group(&"group1".into())
2043 .unwrap()
2044 .edges()
2045 .contains_key(&"key1".into())
2046 );
2047 }
2048
2049 #[test]
2050 fn test_schema_add_group() {
2051 let attribute_schema = AttributeSchema::new(
2052 vec![
2053 (
2054 "key1".into(),
2055 AttributeDataType::new(DataType::Int, AttributeType::Categorical)
2056 .expect("AttributeType was inferred from DataType."),
2057 ),
2058 (
2059 "key2".into(),
2060 AttributeDataType::new(DataType::Float, AttributeType::Continuous)
2061 .expect("AttributeType was inferred from DataType."),
2062 ),
2063 ]
2064 .into_iter()
2065 .collect(),
2066 );
2067
2068 let mut schema = Schema::new_inferred(HashMap::new(), GroupSchema::default());
2069 schema
2070 .add_group(
2071 "group1".into(),
2072 GroupSchema::new(attribute_schema.clone(), AttributeSchema::default()),
2073 )
2074 .unwrap();
2075 assert_eq!(
2076 attribute_schema,
2077 schema.group(&"group1".into()).unwrap().nodes
2078 );
2079
2080 assert!(
2081 schema
2082 .add_group("group1".into(), GroupSchema::default())
2083 .is_err_and(|error| {
2084 matches!(
2085 error,
2086 crate::errors::SchemaError::GroupAlreadyInSchema { .. }
2087 )
2088 })
2089 );
2090 }
2091
2092 #[test]
2093 fn test_schema_remove_group() {
2094 let mut schema = Schema::new_inferred(
2095 vec![("group1".into(), GroupSchema::default())]
2096 .into_iter()
2097 .collect(),
2098 GroupSchema::default(),
2099 );
2100 schema.remove_group(&"group1".into());
2101 assert!(!schema.groups().contains_key(&"group1".into()));
2102 }
2103
2104 #[test]
2105 fn test_schema_freeze_unfreeze() {
2106 let mut schema = Schema::new_inferred(HashMap::new(), GroupSchema::default());
2107 assert_eq!(schema.schema_type(), &SchemaType::Inferred);
2108
2109 schema.freeze();
2110 assert_eq!(schema.schema_type(), &SchemaType::Provided);
2111
2112 schema.unfreeze();
2113 assert_eq!(schema.schema_type(), &SchemaType::Inferred);
2114 }
2115}