Skip to main content

graphrecords_core/graphrecord/
schema.rs

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}