graphannis_core/graph/serialization/
graphml.rs

1use crate::{
2    annostorage::{Match, ValueSearch},
3    errors::{GraphAnnisCoreError, Result},
4    graph::{
5        ANNIS_NS, Graph, NODE_NAME, NODE_NAME_KEY, NODE_TYPE, NODE_TYPE_KEY,
6        update::{GraphUpdate, UpdateEvent},
7    },
8    types::{AnnoKey, Annotation, Component, ComponentType, Edge},
9    util::{join_qname, split_qname},
10};
11use itertools::Itertools;
12use quick_xml::{
13    Reader, Writer,
14    events::{
15        BytesCData, BytesDecl, BytesEnd, BytesStart, BytesText, Event, attributes::Attributes,
16    },
17};
18use std::{
19    cmp::Ordering,
20    collections::{BTreeMap, BTreeSet, HashMap, VecDeque},
21    io::{BufReader, BufWriter, Read, Write},
22    path::{Path, PathBuf},
23    str::FromStr,
24};
25
26/// Import a single GraphML-file as an annotation [`Graph`].
27///
28/// # Returns
29///
30/// A tuple of the graph itself and an optional corpus configuration as string.
31pub fn import<CT: ComponentType, R: Read, F>(
32    input: R,
33    disk_based: bool,
34    progress_callback: F,
35) -> Result<(Graph<CT>, Option<String>)>
36where
37    F: Fn(&str),
38{
39    // Always buffer the read operations
40    let mut input = BufReader::new(input);
41    let mut g = Graph::with_default_graphstorages(disk_based)?;
42    let mut updates = GraphUpdate::default();
43    let mut edge_updates = GraphUpdate::default();
44
45    // read in all nodes and edges, collecting annotation keys on the fly
46    progress_callback("reading GraphML");
47    let config = read_graphml::<CT, BufReader<R>, F>(
48        &mut input,
49        &mut updates,
50        &mut edge_updates,
51        &progress_callback,
52    )?;
53
54    // Append all edges updates after the node updates:
55    // edges would not be added if the nodes they are referring do not exist
56    progress_callback("merging generated events");
57    for event in edge_updates.iter()? {
58        let (_, event) = event?;
59        updates.add_event(event)?;
60    }
61
62    progress_callback("applying imported changes");
63    g.apply_update(&mut updates, &progress_callback)?;
64
65    progress_callback("calculating graph statistics");
66    g.calculate_all_statistics()?;
67
68    for c in g.get_all_components(None, None) {
69        progress_callback(&format!("optimizing implementation for component {}", c));
70        g.optimize_gs_impl(&c)?;
71    }
72
73    Ok((g, config))
74}
75
76/// Read in a single GraphML file from `input` and fill the updates given as
77/// parameters.
78///
79/// In order to create a [`Graph`] from it, you first have to apply the
80/// `node_updates` and then the `edge_updates`. Status updates can be retrieved
81/// by the `progress_updates` closure, that will be called with a message as
82/// argument and can be used e.g. for logging or displaying the status message
83/// to the user.
84///
85/// # Returns
86///
87/// If the GraphML-file contains a corpus configuration, this is returned as a string.
88pub fn read_graphml<CT: ComponentType, R: std::io::BufRead, F: Fn(&str)>(
89    input: &mut R,
90    node_updates: &mut GraphUpdate,
91    edge_updates: &mut GraphUpdate,
92    progress_callback: &F,
93) -> Result<Option<String>> {
94    let mut reader = Reader::from_reader(input);
95    reader.expand_empty_elements(true);
96
97    let mut keys = BTreeMap::new();
98
99    let mut level = 0;
100    let mut in_graph = false;
101    let mut current_node_id: Option<String> = None;
102    let mut current_data_key: Option<String> = None;
103    let mut current_source_id: Option<String> = None;
104    let mut current_target_id: Option<String> = None;
105    let mut current_component: Option<String> = None;
106    let mut current_data_value: Option<String> = None;
107    let mut data: HashMap<AnnoKey, String> = HashMap::new();
108
109    let mut config = None;
110
111    let mut processed_updates = 0;
112
113    let mut buf = Vec::new();
114    loop {
115        match reader.read_event_into(&mut buf)? {
116            Event::Start(ref e) => {
117                level += 1;
118
119                match e.name().0 {
120                    b"graph" if level == 2 => {
121                        in_graph = true;
122                    }
123                    b"key" if level == 2 => {
124                        add_annotation_key(&mut keys, e.attributes())?;
125                    }
126                    b"node" if in_graph && level == 3 => {
127                        data.clear();
128                        // Get the ID of this node
129                        for att in e.attributes() {
130                            let att = att?;
131                            if att.key.0 == b"id" {
132                                current_node_id =
133                                    Some(String::from_utf8_lossy(&att.value).to_string());
134                            }
135                        }
136                    }
137
138                    b"edge" if in_graph && level == 3 => {
139                        data.clear();
140                        // Get the source and target node IDs
141                        for att in e.attributes() {
142                            let att = att?;
143                            if att.key.0 == b"source" {
144                                current_source_id =
145                                    Some(String::from_utf8_lossy(&att.value).to_string());
146                            } else if att.key.0 == b"target" {
147                                current_target_id =
148                                    Some(String::from_utf8_lossy(&att.value).to_string());
149                            } else if att.key.0 == b"label" {
150                                current_component =
151                                    Some(String::from_utf8_lossy(&att.value).to_string());
152                            }
153                        }
154                    }
155
156                    b"data" => {
157                        for att in e.attributes() {
158                            let att = att?;
159                            if att.key.0 == b"key" {
160                                current_data_key =
161                                    Some(String::from_utf8_lossy(&att.value).to_string());
162                            }
163                        }
164                    }
165                    _ => {}
166                }
167            }
168            Event::Text(t) if in_graph && level == 4 && current_data_key.is_some() => {
169                current_data_value = Some(t.unescape()?.to_string());
170            }
171
172            Event::CData(t) => {
173                if let Some(current_data_key) = &current_data_key
174                    && in_graph
175                    && level == 3
176                    && current_data_key == "k0"
177                {
178                    // This is the configuration content
179                    config = Some(String::from_utf8_lossy(&t).to_string());
180                }
181            }
182            Event::End(ref e) => {
183                match e.name().0 {
184                    b"graph" => {
185                        in_graph = false;
186                    }
187                    b"node" => {
188                        add_node(node_updates, &current_node_id, &mut data)?;
189                        current_node_id = None;
190                        processed_updates += 1;
191                        if processed_updates % 1_000_000 == 0 {
192                            progress_callback(&format!(
193                                "Processed {} GraphML nodes and edges",
194                                processed_updates
195                            ));
196                        }
197                    }
198                    b"edge" => {
199                        add_edge::<CT>(
200                            edge_updates,
201                            &current_source_id,
202                            &current_target_id,
203                            &current_component,
204                            &mut data,
205                        )?;
206                        current_source_id = None;
207                        current_target_id = None;
208                        current_component = None;
209                        processed_updates += 1;
210                        if processed_updates % 1_000_000 == 0 {
211                            progress_callback(&format!(
212                                "Processed {} GraphML nodes and edges",
213                                processed_updates
214                            ));
215                        }
216                    }
217                    b"data" => {
218                        if let Some(current_data_key) = current_data_key
219                            && let Some(anno_key) = keys.get(&current_data_key)
220                        {
221                            // Copy all data attributes into our own map
222                            if let Some(v) = current_data_value.take() {
223                                data.insert(anno_key.clone(), v);
224                            } else {
225                                // If there is an end tag without any text
226                                // data event, the value exists but is
227                                // empty.
228                                data.insert(anno_key.clone(), String::default());
229                            }
230                        }
231
232                        current_data_value = None;
233                        current_data_key = None;
234                    }
235                    _ => {}
236                }
237
238                level -= 1;
239            }
240            Event::Eof => {
241                break;
242            }
243            _ => {}
244        }
245        // Clear the buffer after each event
246        buf.clear();
247    }
248    Ok(config)
249}
250
251/// A corpus can consist of several GraphML-files if the corpus is partitioned
252/// and there is a subdirectory with the same name as the basename of the
253/// GraphML file given as argument.
254///
255/// This function finds the files that belong the same corpus for and returns
256/// the paths in the order they should be read. If there is no matching
257/// subdirectory next to the given GraphML-file, the file itself is returned.
258pub fn files_for_corpus<P: AsRef<Path>>(file: P) -> Result<Vec<PathBuf>> {
259    let mut result = Vec::new();
260    if file.as_ref().is_file()
261        && let Some(ext) = file.as_ref().extension()
262        && ext == "graphml"
263    {
264        // Add the root GraphML file first
265        result.push(file.as_ref().to_path_buf());
266
267        // If there is a directory with the same base name as the GraphML file,
268        // search this directory with a BFS for more GraphML-files
269        if let Some(parent_dir) = file.as_ref().parent()
270            && let Some(basename) = file.as_ref().file_stem()
271            && let corpus_dir = parent_dir.join(basename)
272            && corpus_dir.is_dir()
273        {
274            let mut queue = VecDeque::new();
275            queue.push_back(corpus_dir);
276
277            while let Some(current_file) = queue.pop_front() {
278                if current_file.is_dir() {
279                    // Get all files and directories that belong to this parent
280                    // directory and add them to the queue in a predicatble order.
281                    let mut same_level_entries = BTreeMap::new();
282                    for dir_entry in std::fs::read_dir(&current_file)? {
283                        let dir_entry = dir_entry?;
284                        same_level_entries.insert(dir_entry.file_name(), dir_entry.path());
285                    }
286
287                    for p in same_level_entries.into_values() {
288                        queue.push_back(p);
289                    }
290                } else if current_file.is_file()
291                    && let Some(extension) = current_file.extension()
292                    && extension == "graphml"
293                {
294                    result.push(current_file);
295                }
296            }
297        }
298    }
299    Ok(result)
300}
301
302/// Export the GraphML file without any guarantuee on the order of the XML elements.
303///
304/// This is faster than  than [`export_stable_order`].
305pub fn export<CT: ComponentType, W: std::io::Write, F>(
306    graph: &Graph<CT>,
307    graph_configuration: Option<&str>,
308    output: W,
309    progress_callback: F,
310) -> Result<()>
311where
312    F: Fn(&str),
313{
314    // Always buffer the output
315    let output = BufWriter::new(output);
316    let mut writer = Writer::new_with_indent(output, b' ', 4);
317
318    // Add XML declaration
319    let xml_decl = BytesDecl::new("1.0", Some("UTF-8"), None);
320    writer.write_event(Event::Decl(xml_decl))?;
321
322    // Always write the root element
323    writer.write_event(Event::Start(BytesStart::new("graphml")))?;
324
325    // Define all valid annotation ns/name pairs
326    progress_callback("exporting all available annotation keys");
327    let key_id_mapping =
328        write_annotation_keys(graph, graph_configuration.is_some(), false, &mut writer)?;
329
330    // We are writing a single graph
331    let mut graph_start = BytesStart::new("graph");
332    graph_start.push_attribute(("edgedefault", "directed"));
333    // Add parse helper information to allow more efficient parsing
334    graph_start.push_attribute(("parse.order", "nodesfirst"));
335    graph_start.push_attribute(("parse.nodeids", "free"));
336    graph_start.push_attribute(("parse.edgeids", "canonical"));
337
338    writer.write_event(Event::Start(graph_start))?;
339
340    // If graph configuration is given, add it as data element to the graph
341    if let Some(config) = graph_configuration {
342        let mut data_start = BytesStart::new("data");
343        // This is always the first key ID
344        data_start.push_attribute(("key", "k0"));
345        writer.write_event(Event::Start(data_start))?;
346        // Add the annotation value as internal text node
347        writer.write_event(Event::CData(BytesCData::new(config)))?;
348        writer.write_event(Event::End(BytesEnd::new("data")))?;
349    }
350
351    // Write out all nodes
352    progress_callback("exporting nodes");
353    write_nodes(graph, &mut writer, false, &key_id_mapping)?;
354
355    // Write out all edges
356    progress_callback("exporting edges");
357    write_edges(graph, &mut writer, false, &key_id_mapping)?;
358
359    writer.write_event(Event::End(BytesEnd::new("graph")))?;
360    writer.write_event(Event::End(BytesEnd::new("graphml")))?;
361
362    // Make sure to flush the buffered writer
363    writer.into_inner().flush()?;
364
365    Ok(())
366}
367
368/// Export the GraphML file and ensure a stable order of the XML elements.
369///
370/// This is slower than [`export`] but can e.g. be used in tests where the
371/// output should always be the same.
372pub fn export_stable_order<CT: ComponentType, W: std::io::Write, F>(
373    graph: &Graph<CT>,
374    graph_configuration: Option<&str>,
375    output: W,
376    progress_callback: F,
377) -> Result<()>
378where
379    F: Fn(&str),
380{
381    // Always buffer the output
382    let output = BufWriter::new(output);
383    let mut writer = Writer::new_with_indent(output, b' ', 4);
384
385    // Add XML declaration
386    let xml_decl = BytesDecl::new("1.0", Some("UTF-8"), None);
387    writer.write_event(Event::Decl(xml_decl))?;
388
389    // Always write the root element
390    writer.write_event(Event::Start(BytesStart::new("graphml")))?;
391
392    // Define all valid annotation ns/name pairs
393    progress_callback("exporting all available annotation keys");
394    let key_id_mapping =
395        write_annotation_keys(graph, graph_configuration.is_some(), true, &mut writer)?;
396
397    // We are writing a single graph
398    let mut graph_start = BytesStart::new("graph");
399    graph_start.push_attribute(("edgedefault", "directed"));
400    // Add parse helper information to allow more efficient parsing
401    graph_start.push_attribute(("parse.order", "nodesfirst"));
402    graph_start.push_attribute(("parse.nodeids", "free"));
403    graph_start.push_attribute(("parse.edgeids", "canonical"));
404
405    writer.write_event(Event::Start(graph_start))?;
406
407    // If graph configuration is given, add it as data element to the graph
408    if let Some(config) = graph_configuration {
409        let mut data_start = BytesStart::new("data");
410        // This is always the first key ID
411        data_start.push_attribute(("key", "k0"));
412        writer.write_event(Event::Start(data_start))?;
413        // Add the annotation value as internal text node
414        writer.write_event(Event::CData(BytesCData::new(config)))?;
415        writer.write_event(Event::End(BytesEnd::new("data")))?;
416    }
417
418    // Write out all nodes
419    progress_callback("exporting nodes");
420    write_nodes(graph, &mut writer, true, &key_id_mapping)?;
421
422    // Write out all edges
423    progress_callback("exporting edges");
424    write_edges(graph, &mut writer, true, &key_id_mapping)?;
425
426    writer.write_event(Event::End(BytesEnd::new("graph")))?;
427    writer.write_event(Event::End(BytesEnd::new("graphml")))?;
428
429    // Make sure to flush the buffered writer
430    writer.into_inner().flush()?;
431
432    Ok(())
433}
434
435fn write_annotation_keys<CT: ComponentType, W: std::io::Write>(
436    graph: &Graph<CT>,
437    has_graph_configuration: bool,
438    sorted: bool,
439    writer: &mut Writer<W>,
440) -> Result<BTreeMap<AnnoKey, String>> {
441    let mut key_id_mapping = BTreeMap::new();
442    let mut id_counter = 0;
443
444    if has_graph_configuration {
445        let new_id = format!("k{}", id_counter);
446        id_counter += 1;
447
448        let mut key_start = BytesStart::new("key");
449        key_start.push_attribute(("id", new_id.as_str()));
450        key_start.push_attribute(("for", "graph"));
451        key_start.push_attribute(("attr.name", "configuration"));
452        key_start.push_attribute(("attr.type", "string"));
453
454        writer.write_event(Event::Empty(key_start))?;
455    }
456
457    // Create node annotation keys
458    let mut anno_keys = graph.get_node_annos().annotation_keys()?;
459    if sorted {
460        anno_keys.sort_unstable();
461    }
462    for key in anno_keys {
463        if (key.ns != ANNIS_NS || key.name != NODE_NAME) && !key_id_mapping.contains_key(&key) {
464            let new_id = format!("k{}", id_counter);
465            id_counter += 1;
466
467            let qname = join_qname(&key.ns, &key.name);
468
469            let mut key_start = BytesStart::new("key");
470            key_start.push_attribute(("id", new_id.as_str()));
471            key_start.push_attribute(("for", "node"));
472            key_start.push_attribute(("attr.name", qname.as_str()));
473            key_start.push_attribute(("attr.type", "string"));
474
475            writer.write_event(Event::Empty(key_start))?;
476
477            key_id_mapping.insert(key, new_id);
478        }
479    }
480
481    // Create edge annotation keys for all components, but skip auto-generated ones
482    let autogenerated_components: BTreeSet<Component<CT>> =
483        CT::update_graph_index_components(graph)
484            .into_iter()
485            .collect();
486    let mut all_components = graph.get_all_components(None, None);
487    if sorted {
488        all_components.sort_unstable();
489    }
490    for c in all_components {
491        if !autogenerated_components.contains(&c)
492            && let Some(gs) = graph.get_graphstorage(&c)
493        {
494            for key in gs.get_anno_storage().annotation_keys()? {
495                #[allow(clippy::map_entry)]
496                if !key_id_mapping.contains_key(&key) {
497                    let new_id = format!("k{}", id_counter);
498                    id_counter += 1;
499
500                    let qname = join_qname(&key.ns, &key.name);
501
502                    let mut key_start = BytesStart::new("key");
503                    key_start.push_attribute(("id", new_id.as_str()));
504                    key_start.push_attribute(("for", "node"));
505                    key_start.push_attribute(("attr.name", qname.as_str()));
506                    key_start.push_attribute(("attr.type", "string"));
507
508                    writer.write_event(Event::Empty(key_start))?;
509
510                    key_id_mapping.insert(key, new_id);
511                }
512            }
513        }
514    }
515
516    Ok(key_id_mapping)
517}
518
519fn write_data<W: std::io::Write>(
520    anno: Annotation,
521    writer: &mut Writer<W>,
522    key_id_mapping: &BTreeMap<AnnoKey, String>,
523) -> Result<()> {
524    let mut data_start = BytesStart::new("data");
525
526    let key_id = key_id_mapping
527        .get(&anno.key)
528        .ok_or_else(|| GraphAnnisCoreError::GraphMLMissingAnnotationKey(anno.key.clone()))?;
529
530    data_start.push_attribute(("key", key_id.as_str()));
531    writer.write_event(Event::Start(data_start))?;
532    // Add the annotation value as internal text node
533    writer.write_event(Event::Text(BytesText::new(&anno.val)))?;
534    writer.write_event(Event::End(BytesEnd::new("data")))?;
535
536    Ok(())
537}
538
539fn compare_results<T: Ord>(a: &Result<T>, b: &Result<T>) -> Ordering {
540    if let (Ok(a), Ok(b)) = (a, b) {
541        a.cmp(b)
542    } else if a.is_err() {
543        Ordering::Less
544    } else if b.is_err() {
545        Ordering::Greater
546    } else {
547        // Treat two errors as equal
548        Ordering::Equal
549    }
550}
551
552fn write_nodes<CT: ComponentType, W: std::io::Write>(
553    graph: &Graph<CT>,
554    writer: &mut Writer<W>,
555    sorted: bool,
556    key_id_mapping: &BTreeMap<AnnoKey, String>,
557) -> Result<()> {
558    let base_node_iterator =
559        graph
560            .get_node_annos()
561            .exact_anno_search(Some(ANNIS_NS), NODE_TYPE, ValueSearch::Any);
562    let node_iterator: Box<dyn Iterator<Item = Result<Match>>> = if sorted {
563        let it = base_node_iterator.sorted_unstable_by(compare_results);
564        Box::new(it)
565    } else {
566        Box::new(base_node_iterator)
567    };
568
569    for m in node_iterator {
570        let m = m?;
571        let mut node_start = BytesStart::new("node");
572
573        if let Some(id) = graph
574            .get_node_annos()
575            .get_value_for_item(&m.node, &NODE_NAME_KEY)?
576        {
577            node_start.push_attribute(("id", id.as_ref()));
578            let mut node_annotations = graph.get_node_annos().get_annotations_for_item(&m.node)?;
579            if node_annotations.is_empty() {
580                // Write an empty XML element without child nodes
581                writer.write_event(Event::Empty(node_start))?;
582            } else {
583                writer.write_event(Event::Start(node_start))?;
584                // Write all annotations of the node as "data" element, but sort
585                // them using the internal annotation key (k0, k1, k2, etc.)
586                node_annotations.sort_unstable_by_key(|anno| {
587                    key_id_mapping
588                        .get(&anno.key)
589                        .map(|internal_key| internal_key.as_str())
590                        .unwrap_or("")
591                });
592
593                for anno in node_annotations {
594                    if anno.key.ns != ANNIS_NS || anno.key.name != NODE_NAME {
595                        write_data(anno, writer, key_id_mapping)?;
596                    }
597                }
598                writer.write_event(Event::End(BytesEnd::new("node")))?;
599            }
600        }
601    }
602    Ok(())
603}
604
605fn write_edges<CT: ComponentType, W: std::io::Write>(
606    graph: &Graph<CT>,
607    writer: &mut Writer<W>,
608    sorted: bool,
609    key_id_mapping: &BTreeMap<AnnoKey, String>,
610) -> Result<()> {
611    let mut edge_counter = 0;
612    let autogenerated_components: BTreeSet<Component<CT>> =
613        CT::update_graph_index_components(graph)
614            .into_iter()
615            .collect();
616
617    let mut all_components = graph.get_all_components(None, None);
618    if sorted {
619        all_components.sort_unstable();
620    }
621
622    for c in all_components {
623        // Create edge annotation keys for all components, but skip auto-generated ones
624        if !autogenerated_components.contains(&c)
625            && let Some(gs) = graph.get_graphstorage(&c)
626        {
627            let source_nodes_iterator = if sorted {
628                Box::new(gs.source_nodes().sorted_unstable_by(compare_results))
629            } else {
630                gs.source_nodes()
631            };
632            for source in source_nodes_iterator {
633                let source = source?;
634                if let Some(source_id) = graph
635                    .get_node_annos()
636                    .get_value_for_item(&source, &NODE_NAME_KEY)?
637                {
638                    let target_nodes_iterator = if sorted {
639                        Box::new(
640                            gs.get_outgoing_edges(source)
641                                .sorted_unstable_by(compare_results),
642                        )
643                    } else {
644                        gs.get_outgoing_edges(source)
645                    };
646                    for target in target_nodes_iterator {
647                        let target = target?;
648                        if let Some(target_id) = graph
649                            .get_node_annos()
650                            .get_value_for_item(&target, &NODE_NAME_KEY)?
651                        {
652                            let edge = Edge { source, target };
653
654                            let mut edge_id = edge_counter.to_string();
655                            edge_counter += 1;
656                            edge_id.insert(0, 'e');
657
658                            let mut edge_start = BytesStart::new("edge");
659                            edge_start.push_attribute(("id", edge_id.as_str()));
660                            edge_start.push_attribute(("source", source_id.as_ref()));
661                            edge_start.push_attribute(("target", target_id.as_ref()));
662                            // Use the "label" attribute as component type. This is consistent with how Neo4j interprets this non-standard attribute
663                            edge_start.push_attribute(("label", c.to_string().as_ref()));
664
665                            writer.write_event(Event::Start(edge_start))?;
666
667                            // Write all annotations of the node as "data" element, but sort
668                            // them using the internal annotation key (k0, k1, k2, etc.)
669                            let mut edge_annotations =
670                                gs.get_anno_storage().get_annotations_for_item(&edge)?;
671                            edge_annotations.sort_unstable_by_key(|anno| {
672                                key_id_mapping
673                                    .get(&anno.key)
674                                    .map(|internal_key| internal_key.as_str())
675                                    .unwrap_or("")
676                            });
677                            for anno in edge_annotations {
678                                write_data(anno, writer, key_id_mapping)?;
679                            }
680                            writer.write_event(Event::End(BytesEnd::new("edge")))?;
681                        }
682                    }
683                }
684            }
685        }
686    }
687    Ok(())
688}
689
690fn add_annotation_key(keys: &mut BTreeMap<String, AnnoKey>, attributes: Attributes) -> Result<()> {
691    // resolve the ID to the fully qualified annotation name
692    let mut id: Option<String> = None;
693    let mut anno_key: Option<AnnoKey> = None;
694
695    for att in attributes {
696        let att = att?;
697
698        let att_value = String::from_utf8_lossy(&att.value);
699
700        match att.key.0 {
701            b"id" => {
702                id = Some(att_value.to_string());
703            }
704            b"attr.name" => {
705                let (ns, name) = split_qname(att_value.as_ref());
706                anno_key = Some(AnnoKey {
707                    ns: ns.unwrap_or("").into(),
708                    name: name.into(),
709                });
710            }
711            _ => {}
712        }
713    }
714
715    if let (Some(id), Some(anno_key)) = (id, anno_key) {
716        keys.insert(id, anno_key);
717    }
718    Ok(())
719}
720
721fn add_node(
722    node_updates: &mut GraphUpdate,
723    current_node_id: &Option<String>,
724    data: &mut HashMap<AnnoKey, String>,
725) -> Result<()> {
726    if let Some(node_name) = current_node_id {
727        // Insert graph update for node
728        let node_type = data
729            .remove(&NODE_TYPE_KEY)
730            .unwrap_or_else(|| "node".to_string());
731        node_updates.add_event(UpdateEvent::AddNode {
732            node_name: node_name.clone(),
733            node_type,
734        })?;
735        // Add all remaining data entries as annotations
736        for (key, value) in data.drain() {
737            node_updates.add_event(UpdateEvent::AddNodeLabel {
738                node_name: node_name.clone(),
739                anno_ns: key.ns,
740                anno_name: key.name,
741                anno_value: value,
742            })?;
743        }
744    }
745    Ok(())
746}
747
748fn add_edge<CT: ComponentType>(
749    edge_updates: &mut GraphUpdate,
750    current_source_id: &Option<String>,
751    current_target_id: &Option<String>,
752    current_component: &Option<String>,
753    data: &mut HashMap<AnnoKey, String>,
754) -> Result<()> {
755    if let (Some(source), Some(target), Some(component)) =
756        (current_source_id, current_target_id, current_component)
757    {
758        // Insert graph update for this edge
759        if let Ok(component) = Component::<CT>::from_str(component) {
760            edge_updates.add_event(UpdateEvent::AddEdge {
761                source_node: source.clone(),
762                target_node: target.clone(),
763                layer: component.layer.clone(),
764                component_type: component.get_type().to_string(),
765                component_name: component.name.clone(),
766            })?;
767
768            // Add all remaining data entries as annotations
769            for (key, value) in data.drain() {
770                edge_updates.add_event(UpdateEvent::AddEdgeLabel {
771                    source_node: source.clone(),
772                    target_node: target.clone(),
773                    layer: component.layer.clone(),
774                    component_type: component.get_type().to_string(),
775                    component_name: component.name.clone(),
776                    anno_ns: key.ns,
777                    anno_name: key.name,
778                    anno_value: value,
779                })?;
780            }
781        }
782    }
783    Ok(())
784}
785
786#[cfg(test)]
787mod tests {
788    use super::*;
789    use crate::{
790        graph::{DEFAULT_NS, GraphUpdate},
791        types::DefaultComponentType,
792    };
793    use pretty_assertions::assert_eq;
794    use std::borrow::Cow;
795
796    const TEST_CONFIG: &str = r#"[some]
797key = "<value>"
798
799[some.another]
800value = "test""#;
801
802    #[test]
803    fn export_graphml() {
804        // Create a sample graph using the simple type
805        let mut u = GraphUpdate::new();
806        u.add_event(UpdateEvent::AddNode {
807            node_name: "first_node".to_string(),
808            node_type: "node".to_string(),
809        })
810        .unwrap();
811        u.add_event(UpdateEvent::AddNode {
812            node_name: "second_node".to_string(),
813            node_type: "node".to_string(),
814        })
815        .unwrap();
816        u.add_event(UpdateEvent::AddNodeLabel {
817            node_name: "first_node".to_string(),
818            anno_ns: DEFAULT_NS.to_string(),
819            anno_name: "an_annotation".to_string(),
820            anno_value: "something <strong>important</strong>".to_string(),
821        })
822        .unwrap();
823
824        u.add_event(UpdateEvent::AddEdge {
825            source_node: "first_node".to_string(),
826            target_node: "second_node".to_string(),
827            component_type: "Edge".to_string(),
828            layer: "some_ns".to_string(),
829            component_name: "test_component".to_string(),
830        })
831        .unwrap();
832
833        let mut g: Graph<DefaultComponentType> = Graph::new(false).unwrap();
834        g.apply_update(&mut u, |_| {}).unwrap();
835
836        // export to GraphML, read generated XML and compare it
837        let mut xml_data: Vec<u8> = Vec::default();
838        export(&g, Some(TEST_CONFIG), &mut xml_data, |_| {}).unwrap();
839        let expected = include_str!("graphml_example.graphml");
840        let actual = String::from_utf8(xml_data).unwrap();
841        assert_eq!(expected, actual);
842    }
843
844    #[test]
845    fn export_graphml_sorted() {
846        // Create a sample graph using the simple type
847        let mut u = GraphUpdate::new();
848
849        u.add_event(UpdateEvent::AddNode {
850            node_name: "1".to_string(),
851            node_type: "node".to_string(),
852        })
853        .unwrap();
854        u.add_event(UpdateEvent::AddNode {
855            node_name: "2".to_string(),
856            node_type: "node".to_string(),
857        })
858        .unwrap();
859        u.add_event(UpdateEvent::AddNodeLabel {
860            node_name: "1".to_string(),
861            anno_ns: DEFAULT_NS.to_string(),
862            anno_name: "an_annotation".to_string(),
863            anno_value: "something".to_string(),
864        })
865        .unwrap();
866
867        u.add_event(UpdateEvent::AddEdge {
868            source_node: "1".to_string(),
869            target_node: "2".to_string(),
870            component_type: "Edge".to_string(),
871            layer: "some_ns".to_string(),
872            component_name: "test_component".to_string(),
873        })
874        .unwrap();
875
876        let mut g: Graph<DefaultComponentType> = Graph::new(false).unwrap();
877        g.apply_update(&mut u, |_| {}).unwrap();
878
879        // export to GraphML, read generated XML and compare it
880        let mut xml_data: Vec<u8> = Vec::default();
881        export_stable_order(&g, Some(TEST_CONFIG), &mut xml_data, |_| {}).unwrap();
882        let expected = include_str!("graphml_example sorted.graphml");
883        let actual = String::from_utf8(xml_data).unwrap();
884        assert_eq!(expected, actual);
885    }
886
887    #[test]
888    fn import_graphml() {
889        let input_xml = std::io::Cursor::new(
890            include_str!("graphml_example.graphml")
891                .as_bytes()
892                .to_owned(),
893        );
894        let (g, config_str) = import(input_xml, false, |_| {}).unwrap();
895
896        // Check that all nodes, edges and annotations have been created
897        let first_node_id = g
898            .node_annos
899            .get_node_id_from_name("first_node")
900            .unwrap()
901            .unwrap();
902        let second_node_id = g
903            .node_annos
904            .get_node_id_from_name("second_node")
905            .unwrap()
906            .unwrap();
907
908        let first_node_annos = g
909            .get_node_annos()
910            .get_annotations_for_item(&first_node_id)
911            .unwrap();
912        assert_eq!(3, first_node_annos.len());
913        assert_eq!(
914            Some(Cow::Borrowed("something <strong>important</strong>")),
915            g.get_node_annos()
916                .get_value_for_item(
917                    &first_node_id,
918                    &AnnoKey {
919                        ns: DEFAULT_NS.into(),
920                        name: "an_annotation".into(),
921                    }
922                )
923                .unwrap()
924        );
925
926        assert_eq!(
927            2,
928            g.get_node_annos()
929                .get_annotations_for_item(&second_node_id)
930                .unwrap()
931                .len()
932        );
933
934        let component = g.get_all_components(Some(DefaultComponentType::Edge), None);
935        assert_eq!(1, component.len());
936        assert_eq!("some_ns", component[0].layer);
937        assert_eq!("test_component", component[0].name);
938
939        let test_gs = g.get_graphstorage_as_ref(&component[0]).unwrap();
940        assert_eq!(
941            Some(1),
942            test_gs.distance(first_node_id, second_node_id).unwrap()
943        );
944
945        assert_eq!(Some(TEST_CONFIG), config_str.as_deref());
946    }
947
948    #[test]
949    fn test_partitioned_file_import_order() {
950        let example_corpus = Path::new("tests/partioned-graphml/single_sentence.graphml");
951        assert!(example_corpus.is_file());
952
953        let result = files_for_corpus(example_corpus).unwrap();
954        assert_eq!(3, result.len());
955
956        assert_eq!(
957            vec!["tests", "partioned-graphml", "single_sentence.graphml"],
958            result[0].components().map(|c| c.as_os_str()).collect_vec()
959        );
960        assert_eq!(
961            vec![
962                "tests",
963                "partioned-graphml",
964                "single_sentence",
965                "zossen.graphml"
966            ],
967            result[1].components().map(|c| c.as_os_str()).collect_vec()
968        );
969        assert_eq!(
970            vec![
971                "tests",
972                "partioned-graphml",
973                "single_sentence",
974                "subcorpus1",
975                "anotherdocument.graphml"
976            ],
977            result[2].components().map(|c| c.as_os_str()).collect_vec()
978        );
979    }
980
981    #[test]
982    fn test_non_partitioned_file_import_order() {
983        let example_corpus = Path::new("tests/single_sentence.graphml");
984        assert!(example_corpus.is_file());
985
986        let result = files_for_corpus(example_corpus).unwrap();
987        assert_eq!(1, result.len());
988        assert_eq!(
989            vec!["tests", "single_sentence.graphml"],
990            result[0].components().map(|c| c.as_os_str()).collect_vec()
991        );
992    }
993}
graphannis_core/graph/serialization/graphml.rs

graphannis_core/graph/serialization/
graphml.rs
