move prefixing with subgraph to neo4j creation stage

Author: jamesamcl

02_assign_ids/grebi_assign_ids/src/main.rs

@@ -118,15 +118,18 @@ fn main() {
             }
         }
 
+        writer.write_all("{\"grebi:nodeId\":\"".as_bytes()).unwrap();
+
         let group = id_to_group.get(id.unwrap());
-        if !group.is_some() {
-            panic!("could not find identifier group for id: {}", String::from_utf8(id.unwrap().to_vec()).unwrap());
+        if group.is_some() {
+            writer.write_all(group.unwrap().as_slice()).unwrap();
+        } else {
+            writer.write_all(id.unwrap()).unwrap();
         }
 
-        writer.write_all("{\"grebi:nodeId\":\"".as_bytes()).unwrap();
-        writer.write_all(group.unwrap().as_slice()).unwrap();
         writer.write_all("\"".as_bytes()).unwrap();
 
+
         json.rewind();
         while json.peek().kind != JsonTokenType::EndObject {
 

02_assign_ids/grebi_identifiers2groups/src/main.rs

@@ -18,20 +18,15 @@ struct Args {
     #[arg(long)]
     add_group: Vec<String>,
 
-    #[arg(long)]
-    add_prefix: String, // used to prepend the subgraph name like hra_kg:g:
-
-    #[arg(long)]
-    prealloc_size: usize
 }
 
 
 fn main() {
 
 	let args = Args::parse();
 
-	let mut group_to_entities:HashMap<u64, HashSet<Vec<u8>>> = HashMap::with_capacity(args.prealloc_size);
-	let mut entity_to_group:HashMap<Vec<u8>, u64> = HashMap::with_capacity(args.prealloc_size);
+	let mut group_to_entities:HashMap<u64, HashSet<Vec<u8>>> = HashMap::new();
+	let mut entity_to_group:HashMap<Vec<u8>, u64> = HashMap::new();
 
 	let mut next_group_id:u64 = 1;
 
@@ -127,6 +122,12 @@ fn main() {
 
 	for group in group_to_entities {
 
+		if group.1.len() == 1 {
+			// this is a unique id with no equivalences, no need to
+			// write it as a group.
+			continue;
+		}
+
 		n2 = n2 + 1;
 
 		// writer.write_all("group_".as_bytes()).unwrap();
@@ -140,7 +141,6 @@ fn main() {
 
 		for entity in sorted_ids {
 			if is_first_value {
-				writer.write_all(&args.add_prefix.as_bytes()).unwrap();
 				writer.write_all(entity.as_slice()).unwrap();
 				writer.write_all("\t".as_bytes()).unwrap();
 				is_first_value = false;

06_prepare_db_import/grebi_make_neo_csv/src/main.rs

@@ -40,12 +40,17 @@ struct Args {
 
     #[arg(long)]
     out_id_edges_csv_path: String,
+
+    #[arg(long)]
+    add_prefix: String, // used to prepend the subgraph name like hra_kg:g:
 }
 
 fn main() -> std::io::Result<()> {
 
     let args = Args::parse();
 
+    let add_prefix = args.add_prefix.as_bytes();
+
     let start_time = std::time::Instant::now();
 
 
@@ -123,7 +128,7 @@ fn main() -> std::io::Result<()> {
 
         let sliced = SlicedEntity::from_json(&line);
 
-        write_node(&line, &sliced, &all_entity_props, &mut nodes_writer, &mut id_edges_writer);
+        write_node(&line, &sliced, &all_entity_props, &mut nodes_writer, &mut id_edges_writer, &add_prefix);
 
         n_nodes = n_nodes + 1;
         if n_nodes % 1000000 == 0 {
@@ -146,7 +151,7 @@ fn main() -> std::io::Result<()> {
 
         let sliced = SlicedEdge::from_json(&line);
 
-        write_edge(&line, sliced, &all_edge_props, &mut edges_writer);
+        write_edge(&line, sliced, &all_edge_props, &mut edges_writer, &add_prefix);
 
         n_edges = n_edges + 1;
         if n_edges % 1000000 == 0 {
@@ -165,12 +170,13 @@ fn main() -> std::io::Result<()> {
     Ok(())
 }
 
-fn write_node(src_line:&[u8], entity:&SlicedEntity, all_node_props:&HashSet<String>, nodes_writer:&mut BufWriter<&File>, id_edges_writer:&mut BufWriter<&File>) {
+fn write_node(src_line:&[u8], entity:&SlicedEntity, all_node_props:&HashSet<String>, nodes_writer:&mut BufWriter<&File>, id_edges_writer:&mut BufWriter<&File>, add_prefix:&[u8]) {
 
     let refs:Map<String,Value> = serde_json::from_slice(entity._refs.unwrap()).unwrap();
 
     // grebi:nodeId
     nodes_writer.write_all(b"\"").unwrap();
+    nodes_writer.write_all(&add_prefix).unwrap();
     write_escaped_value(entity.id, nodes_writer);
     nodes_writer.write_all(b"\",\"").unwrap();
 
@@ -215,7 +221,7 @@ fn write_node(src_line:&[u8], entity:&SlicedEntity, all_node_props:&HashSet<Stri
                 if header_prop.as_bytes() == row_prop.key {
                     if row_prop.key == "id".as_bytes() {
                         for val in row_prop.values.iter() {
-                            write_id_row(val, id_edges_writer, &entity.id);
+                            write_id_row(val, id_edges_writer, &entity.id, &add_prefix);
                         }
                     }
                     for val in row_prop.values.iter() {
@@ -246,17 +252,20 @@ fn write_node(src_line:&[u8], entity:&SlicedEntity, all_node_props:&HashSet<Stri
     nodes_writer.write_all(b"\n").unwrap();
 }
 
-fn write_edge(src_line:&[u8], edge:SlicedEdge, all_edge_props:&HashSet<String>, edges_writer: &mut BufWriter<&File>) {
+fn write_edge(src_line:&[u8], edge:SlicedEdge, all_edge_props:&HashSet<String>, edges_writer: &mut BufWriter<&File>, add_prefix:&[u8]) {
 
     let refs:Map<String,Value> = serde_json::from_slice(edge._refs.unwrap()).unwrap();
 
     edges_writer.write_all(b"\"").unwrap();
+    write_escaped_value(&add_prefix, edges_writer);
     write_escaped_value(edge.from, edges_writer);
     edges_writer.write_all(b"\",\"").unwrap();
     write_escaped_value(edge.edge_type, edges_writer);
     edges_writer.write_all(b"\",\"").unwrap();
+    write_escaped_value(&add_prefix, edges_writer);
     write_escaped_value(edge.to, edges_writer);
     edges_writer.write_all(b"\",\"").unwrap();
+    write_escaped_value(&add_prefix, edges_writer);
     write_escaped_value(edge.edge_id, edges_writer);
     edges_writer.write_all(b"\",\"").unwrap();
 
@@ -340,7 +349,7 @@ fn parse_json_and_write(buf:&[u8], refs:&Map<String,Value>, writer:&mut BufWrite
     }
 }
 
-fn write_id_row(val:&SlicedPropertyValue, id_edges_writer:&mut BufWriter<&File>, grebi_node_id:&[u8]) {
+fn write_id_row(val:&SlicedPropertyValue, id_edges_writer:&mut BufWriter<&File>, grebi_node_id:&[u8], add_prefix:&[u8]) {
 
     let actual_id = {
         if val.kind == JsonTokenType::StartObject {
@@ -356,6 +365,7 @@ fn write_id_row(val:&SlicedPropertyValue, id_edges_writer:&mut BufWriter<&File>,
     };
 
     id_edges_writer.write_all(b"\"").unwrap();
+    write_escaped_value(&add_prefix, id_edges_writer);
     write_escaped_value(grebi_node_id, id_edges_writer);
     id_edges_writer.write_all(b"\",\"").unwrap();
     write_escaped_value(b"id", id_edges_writer);

nextflow/01_create_subgraph.nf

@@ -10,7 +10,6 @@ params.config = "$GREBI_CONFIG"
 params.subgraph = "$GREBI_SUBGRAPH"
 params.timestamp = "$GREBI_TIMESTAMP"
 params.is_ebi = "$GREBI_IS_EBI"
-params.max_entities = "$GREBI_MAX_ENTITIES"
 
 workflow {
 
@@ -128,8 +127,6 @@ process build_equiv_groups {
     set -Eeuo pipefail
     cat ${identifiers_tsv} \
         | ${params.home}/target/release/grebi_identifiers2groups \
-            --add-prefix ${params.subgraph}:g: \
-            --prealloc-size ${params.max_entities} \
             ${buildAddEquivGroupArgs(additional_equivalence_groups)} \
         > groups.txt
     """
@@ -302,7 +299,8 @@ process prepare_neo {
       --in-edges-jsonl ${edges_jsonl} \
       --out-nodes-csv-path neo_nodes_${params.subgraph}_${task.index}.csv \
       --out-edges-csv-path neo_edges_${params.subgraph}_${task.index}.csv \
-      --out-id-edges-csv-path neo_edges_ids_${params.subgraph}_${task.index}.csv
+      --out-id-edges-csv-path neo_edges_ids_${params.subgraph}_${task.index}.csv \
+      --add-prefix ${params.subgraph}:
     """
 }
 
@@ -406,7 +404,7 @@ process create_solr_nodes_core {
 process create_solr_edges_core {
     cache "lenient"
     memory "64 GB" 
-    time "3d"
+    time "23h"
     cpus "16"
 
     publishDir "${params.tmp}/${params.config}/${params.subgraph}/solr_cores", overwrite: true, saveAs: { filename -> filename.replace("solr/data/", "") }