perf(index): lightweight structure

Cargo.toml

@@ -6,10 +6,12 @@
   rio_api = '0.8.4'
   rio_turtle = '0.8.4'
   rstest = '0.21.0'
+serde_json = "1.0.127"
   serde_yml = '0.0.10'
   slog = '2.7.0'
   slog-async = '2.8.0'
   slog-term = '2.9.0'
+smallvec = { version = "1.13.2", features = ["serde"] }
   tempfile = '3.10.1'
 
   [dependencies.clap]

src/index.rs

@@ -1,38 +1,141 @@
 use rio_api::parser::TriplesParser;
 use rio_turtle::TurtleError;
-use std::{io::Write, path::Path};
+use serde::{Deserialize, Serialize};
+use smallvec::{smallvec, SmallVec};
+use std::{
+    collections::HashMap,
+    hash::{DefaultHasher, Hash, Hasher},
+    path::Path,
+};
 
 use crate::{
     io,
     rdf_types::{Triple, TripleView},
 };
 
-fn index_triple(t: Triple, out: &mut impl Write) {
+/// Stores a mapping from hashed instance uri to their types
+#[derive(Serialize, Deserialize)]
+pub struct TypeIndex {
+    pub types: Vec<String>,
+    map: HashMap<u64, SmallVec<[usize; 1]>>,
+}
+
+impl TypeIndex {
+    fn hash(&mut self, s: &str) -> u64 {
+        let mut hasher = DefaultHasher::new();
+        s.hash(&mut hasher);
+        hasher.finish().to_le()
+    }
+
+    pub fn from_iter(type_map: impl Iterator<Item = (String, String)>) -> Self {
+        let vec: Vec<(String, String)> = type_map.collect();
+        let mut idx = TypeIndex::new();
+        idx.types = vec
+            .clone()
+            .iter()
+            .map(|(_, t)| t.clone())
+            .collect::<std::collections::HashSet<String>>()
+            .into_iter()
+            .collect();
+
+        vec.iter().for_each(|(subject, type_uri)| {
+            idx.insert(&subject.to_string(), &type_uri.to_string())
+                .unwrap()
+        });
+
+        return idx;
+    }
+
+    pub fn new() -> Self {
+        TypeIndex {
+            types: Vec::new(),
+            map: HashMap::new(),
+        }
+    }
+
+    // Insert input subject-type mapping into the index.
+    // The index will store the hash of the subject.
+    pub fn insert(&mut self, subject_uri: &str, type_uri: &str) -> Result<(), std::io::Error> {
+        let key = self.hash(subject_uri);
+        let type_idx: usize;
+
+        // Get type index or add a new one
+        if self.types.contains(&type_uri.to_string()) {
+            type_idx = self.types.iter().position(|x| *x == type_uri).unwrap();
+        } else {
+            type_idx = self.types.len();
+            self.types.push(type_uri.to_string());
+        }
+        // Insert mapping into the index
+        match self.map.get_mut(&key) {
+            Some(v) => {
+                v.push(type_idx);
+            }
+            None => {
+                self.map.insert(key, smallvec![type_idx]);
+            }
+        }
+
+        Ok(())
+    }
+
+    pub fn get(&mut self, subject_key: &str) -> Option<Vec<&String>> {
+        let key = self.hash(subject_key);
+        self.map
+            .get(&key)
+            .map(|v| v.iter().map(|i| &self.types[*i]).collect())
+    }
+}
+
+fn index_triple(t: Triple, index: &mut TypeIndex) {
     if t.predicate.iri.as_str() == "http://www.w3.org/1999/02/22-rdf-syntax-ns#type" {
-        let r = || -> std::io::Result<()> {
-            out.write_all(t.to_string().as_bytes())?;
-            out.write_all(b" .\n")
-        }();
+        let r = { index.insert(&t.subject.to_string(), &t.object.to_string()) };
 
         if let Err(e) = r {
             panic!("Error writting to out buffer: {e}");
         }
     }
 }
 
-pub fn create_type_map(input: &Path, output: &Path) {
+pub fn create_type_index(input: &Path, output: &Path) {
     let buf_in = io::get_reader(input);
-    let mut buf_out = io::get_writer(output);
+    let buf_out = io::get_writer(output);
     let mut triples = io::parse_ntriples(buf_in);
+    let mut index = TypeIndex::new();
 
     while !triples.is_end() {
         let _ = triples
             .parse_step(&mut |t: TripleView| {
-                index_triple(t.into(), &mut buf_out);
+                index_triple(t.into(), &mut index);
                 Result::<(), TurtleError>::Ok(())
             })
             .inspect_err(|e| {
                 panic!("Parsing error occured: {e}");
             });
     }
+    let _ = serde_json::to_writer(buf_out, &index);
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    // Test the parsing of a triple.
+    fn index_from_iter() {
+        let vals = vec![
+            ("urn:Alice", "urn:Person"),
+            ("urn:Alice", "urn:Employee"),
+            ("urn:ACME", "urn:Organization"),
+        ]
+        .into_iter()
+        .map(|(a, b)| (a.to_string(), b.to_string()));
+
+        let mut idx = TypeIndex::from_iter(vals);
+
+        assert_eq!(
+            idx.get("urn:Alice").unwrap(),
+            vec!["urn:Person", "urn:Employee"]
+        );
+        println!("{}", serde_json::to_string(&idx).unwrap());
+    }
 }

src/io.rs

@@ -1,4 +1,4 @@
-use crate::rules::Rules;
+use crate::{index::TypeIndex, rules::Rules};
 use rio_turtle::NTriplesParser;
 use std::{
     fs::File,
@@ -46,7 +46,15 @@ pub fn parse_ntriples(reader: impl BufRead) -> NTriplesParser<impl BufRead> {
 pub fn parse_rules(path: &Path) -> Rules {
     return match File::open(path) {
         Ok(file) => serde_yml::from_reader(file).expect("Error parsing rules file."),
-        Err(e) => panic!("Cannot open file '{:?}': '{}'.", path, e),
+        Err(e) => panic!("Cannot open rules file '{:?}': '{}'.", path, e),
+    };
+}
+
+// Parse yaml type index
+pub fn parse_index(path: &Path) -> TypeIndex {
+    return match File::open(path) {
+        Ok(file) => serde_json::from_reader(file).expect("Error parsing index file."),
+        Err(e) => panic!("Cannot open index file '{:?}': '{}'.", path, e),
     };
 }
 

src/log.rs

@@ -20,7 +20,7 @@ pub fn create_logger(use_stdout: bool) -> Arc<Logger> {
         .fuse();
 
     let drain = slog_async::Async::new(drain)
-        .chan_size(5_000_000)
+        .chan_size(1_000)
         .build()
         .fuse();
 

src/main.rs

@@ -10,7 +10,7 @@ mod rules;
 
 // Define the imports.
 use crate::{
-    index::create_type_map,
+    index::create_type_index,
     log::{create_logger, info},
     pseudo::pseudonymize_graph,
 };
@@ -87,7 +87,7 @@ fn main() {
     match cli.command {
         Subcommands::Index(args) => {
             info!(log, "Args: {:?}", args);
-            create_type_map(&args.input, &args.output)
+            create_type_index(&args.input, &args.output)
         }
         Subcommands::Pseudo(args) => {
             info!(log, "Args: {:?}", args);

src/pseudo.rs

@@ -1,13 +1,13 @@
 use rio_api::parser::TriplesParser;
 use rio_turtle::TurtleError;
 use std::{
-    collections::HashMap,
-    io::{BufRead, Write},
+    io::Write,
     path::{Path, PathBuf},
 };
 
 use crate::{
     crypto::{new_pseudonymizer, Pseudonymize},
+    index::TypeIndex,
     io,
     log::Logger,
     rdf_types::*,
@@ -19,7 +19,7 @@ use crate::{
 fn process_triple(
     triple: Triple,
     rules_config: &Rules,
-    node_to_type: &HashMap<String, String>,
+    node_to_type: &mut TypeIndex,
     out: &mut impl Write,
     hasher: &dyn Pseudonymize,
 ) {
@@ -35,24 +35,6 @@ fn process_triple(
     }
 }
 
-// Create a index mapping node -> type from an input ntriples buffer
-fn load_type_map(input: impl BufRead) -> HashMap<String, String> {
-    let mut node_to_type: HashMap<String, String> = HashMap::new();
-    let mut triples = io::parse_ntriples(input);
-
-    while !triples.is_end() {
-        let _: Result<(), TurtleError> = triples.parse_step(&mut |t| {
-            node_to_type.insert(
-                t.subject.to_string().replace(['<', '>'], ""),
-                t.object.to_string().replace(['<', '>'], ""),
-            );
-            Ok(())
-        });
-    }
-
-    return node_to_type;
-}
-
 pub fn pseudonymize_graph(
     _: &Logger,
     input: &Path,
@@ -62,11 +44,10 @@ pub fn pseudonymize_graph(
     secret_path: &Option<PathBuf>,
 ) {
     let buf_input = io::get_reader(input);
-    let buf_index = io::get_reader(index_path);
     let mut buf_output = io::get_writer(output);
 
     let rules = io::parse_rules(rules_path);
-    let node_to_type: HashMap<String, String> = load_type_map(buf_index);
+    let mut type_index = io::parse_index(index_path);
 
     let secret = secret_path.as_ref().map(io::read_bytes);
     let pseudonymizer = new_pseudonymizer(None, secret);
@@ -80,7 +61,7 @@ pub fn pseudonymize_graph(
                 process_triple(
                     t.into(),
                     &rules,
-                    &node_to_type,
+                    &mut type_index,
                     &mut buf_output,
                     &pseudonymizer,
                 );
@@ -102,14 +83,14 @@ mod tests {
 
     #[test]
     // Test the parsing of a triple.
-    fn encrypt_nt_file() {
+    fn pseudo_nt_file() {
         let logger = log::create_logger(true);
 
         let dir = tempdir().unwrap();
         let input_path = Path::new("tests/data/test.nt");
         let rules_path = Path::new("tests/data/rules.yaml");
         let output_path = dir.path().join("output.nt");
-        let type_map_path = Path::new("tests/data/type_map.nt");
+        let type_map_path = Path::new("tests/data/type_index.json");
         let key = None;
         pseudonymize_graph(
             &logger,