Node merge

Author: vish-cs

pipeline/ingestion/src/main/java/org/datacommons/ingestion/data/CacheReader.java

@@ -141,7 +141,7 @@ public NodesEdges parseArcRow(String row, Counter mcfNodesWithoutTypeCounter) {
                   Node.builder()
                       .subjectId(nodeId)
                       .value(nodeValue)
-                      .bytes(bytes)
+                      .bytes(bytes != null ? bytes.toByteArray() : new byte[0])
                       .name(entity.getName())
                       .types(types)
                       .build());

pipeline/ingestion/src/main/java/org/datacommons/ingestion/data/Edge.java

@@ -14,6 +14,9 @@ public class Edge implements Serializable {
   private String objectId;
   private String provenance;
 
+  @SuppressWarnings("unused")
+  private Edge() {}
+
   // Private constructor to enforce use of Builder
   private Edge(Builder builder) {
     this.subjectId = builder.subjectId;

pipeline/ingestion/src/main/java/org/datacommons/ingestion/data/GraphReader.java

@@ -1,6 +1,5 @@
 package org.datacommons.ingestion.data;
 
-import com.google.cloud.ByteArray;
 import com.google.cloud.spanner.Mutation;
 import java.io.IOException;
 import java.io.Serializable;
@@ -9,8 +8,10 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.transforms.Combine;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.Flatten;
@@ -40,6 +41,118 @@ public class GraphReader implements Serializable {
   private static final String DATCOM_AGGREGATE = "DataCommonsAggregate";
   private static final String IMPORT_METADATA_FILE = "import_metadata_mcf.mcf";
 
+  public static PCollection<Node> combineNodes(PCollection<Node> nodes) {
+    return nodes
+        .apply(
+            "MapNodesToKV",
+            ParDo.of(
+                new DoFn<Node, KV<String, Node>>() {
+                  @ProcessElement
+                  public void processElement(
+                      @Element Node node, OutputReceiver<KV<String, Node>> receiver) {
+                    receiver.output(KV.of(node.getSubjectId(), node));
+                  }
+                }))
+        .apply(
+            "CombineNodes",
+            Combine.perKey(
+                new Combine.CombineFn<Node, List<Node>, Node>() {
+                  @Override
+                  public List<Node> createAccumulator() {
+                    return new ArrayList<>();
+                  }
+
+                  @Override
+                  public List<Node> addInput(List<Node> accumulator, Node input) {
+                    accumulator.add(input);
+                    return accumulator;
+                  }
+
+                  @Override
+                  public List<Node> mergeAccumulators(Iterable<List<Node>> accumulators) {
+                    List<Node> merged = new ArrayList<>();
+                    for (List<Node> acc : accumulators) {
+                      merged.addAll(acc);
+                    }
+                    return merged;
+                  }
+
+                  @Override
+                  public Node extractOutput(List<Node> accumulator) {
+                    if (accumulator.isEmpty()) return null;
+                    Node first = accumulator.get(0);
+                    Node.Builder builder =
+                        Node.builder()
+                            .subjectId(first.getSubjectId())
+                            .value(first.getValue())
+                            .name(first.getName())
+                            .types(first.getTypes())
+                            .bytes(first.getBytes());
+
+                    Set<String> types = new java.util.TreeSet<>();
+                    for (Node n : accumulator) {
+                      types.addAll(n.getTypes());
+                      if (!n.getValue().isEmpty()) {
+                        builder.value(n.getValue());
+                      }
+                      if (!n.getName().isEmpty()) {
+                        builder.name(n.getName());
+                      }
+                      if (n.getBytes().length > 0) {
+                        builder.bytes(n.getBytes());
+                      }
+                    }
+                    if (types.size() > 1 && types.contains("ProvisionalNode")) {
+                      types.remove("ProvisionalNode");
+                    }
+                    builder.types(new ArrayList<>(types));
+                    return builder.build();
+                  }
+                }))
+        .apply(
+            "ExtractNodes",
+            ParDo.of(
+                new DoFn<KV<String, Node>, Node>() {
+                  @ProcessElement
+                  public void processElement(
+                      @Element KV<String, Node> element, OutputReceiver<Node> receiver) {
+                    receiver.output(element.getValue());
+                  }
+                }));
+  }
+
+  public static PCollection<Mutation> nodeToMutations(
+      PCollection<Node> nodes, SpannerClient spannerClient) {
+    return nodes.apply(
+        "NodesToMutations",
+        ParDo.of(
+            new DoFn<Node, Mutation>() {
+              @ProcessElement
+              public void processElement(@Element Node node, OutputReceiver<Mutation> receiver) {
+                Mutation mutation = spannerClient.toNodeMutation(node);
+                if (mutation != null) {
+                  receiver.output(mutation);
+                }
+              }
+            }));
+  }
+
+  public static PCollection<Mutation> edgeToMutations(
+      PCollection<Edge> edges, SpannerClient spannerClient) {
+    return edges.apply(
+        "EdgesToMutations",
+        ParDo.of(
+            new DoFn<Edge, Mutation>() {
+              @ProcessElement
+              public void processElement(@Element Edge edge, OutputReceiver<Mutation> receiver) {
+                Mutation mutation = spannerClient.toEdgeMutation(edge);
+                if (mutation != null) {
+                  receiver.output(mutation);
+                }
+              }
+            }));
+  }
+
   public static List<Node> graphToNodes(McfGraph graph, Counter mcfNodesWithoutTypeCounter) {
     List<Node> nodes = new ArrayList<>();
     for (Map.Entry<String, PropertyValues> nodeEntry : graph.getNodesMap().entrySet()) {
@@ -79,10 +192,11 @@ public static List<Node> graphToNodes(McfGraph graph, Counter mcfNodesWithoutTyp
               node = Node.builder();
               node.subjectId(PipelineUtils.generateObjectValueKey(val.getValue()));
               if (PipelineUtils.storeValueAsBytes(entry.getKey())) {
-                node.bytes(ByteArray.copyFrom(PipelineUtils.compressString(val.getValue())));
+                node.bytes(PipelineUtils.compressString(val.getValue()));
               } else {
                 node.value(val.getValue());
               }
+              node.types(List.of(ValueType.TEXT.toString()));
               nodes.add(node.build());
             }
           }
@@ -101,13 +215,13 @@ public static PCollection<McfGraph> getProvenanceMcf(
     String defaultProvenance =
         "Node: dcid:dc/base/" + importName + "\n" + "typeOf: dcid:Provenance\n";
     mcfList.add(GraphUtils.convertToGraph(defaultProvenance));
+    // try {
+    //   mcfList.add(GraphUtils.convertToGraph(PipelineUtils.getGCSFileContent(metadataFile)));
+    // } catch (IOException e) {
+    //   LOGGER.warn("Failed to read provenance metadata file: " + e.getMessage());
+    // }
     try {
-      mcfList.addAll(GraphUtils.readMcfString(PipelineUtils.getGcsFileContent(metadataFile)));
-    } catch (IOException e) {
-      LOGGER.warn("Failed to read provenance metadata file: " + e.getMessage());
-    }
-    try {
-      mcfList.addAll(GraphUtils.readMcfString(PipelineUtils.getGcsFileContent(provenanceFile)));
+      mcfList.add(GraphUtils.convertToGraph(PipelineUtils.getGCSFileContent(provenanceFile)));
     } catch (IOException e) {
       LOGGER.warn("Failed to read provenance metadata file: " + e.getMessage());
     }
@@ -214,6 +328,42 @@ public void processElement(
             }));
   }
 
+  public static PCollection<Node> mcfToNodes(
+      PCollection<McfGraph> graph, Counter nodeCounter, Counter mcfNodesWithoutTypeCounter) {
+    return graph.apply(
+        "McfToNodes",
+        ParDo.of(
+            new DoFn<McfGraph, Node>() {
+              @ProcessElement
+              public void processElement(@Element McfGraph element, OutputReceiver<Node> receiver) {
+                List<Node> nodes = graphToNodes(element, mcfNodesWithoutTypeCounter);
+                for (Node node : nodes) {
+                  // LOGGER.info("Node: {}", node.toString());
+                  receiver.output(node);
+                }
+                nodeCounter.inc(nodes.size());
+              }
+            }));
+  }
+
+  public static PCollection<Edge> mcfToEdges(
+      PCollection<McfGraph> graph, String provenance, Counter edgeCounter) {
+    return graph.apply(
+        "McfToEdges",
+        ParDo.of(
+            new DoFn<McfGraph, Edge>() {
+              @ProcessElement
+              public void processElement(@Element McfGraph element, OutputReceiver<Edge> receiver) {
+                List<Edge> edges = graphToEdges(element, provenance);
+                for (Edge edge : edges) {
+                  receiver.output(edge);
+                  // LOGGER.info("Edge : {}", edge.toString());
+                }
+                edgeCounter.inc(edges.size());
+              }
+            }));
+  }
+
   public static PCollection<KV<String, Mutation>> graphToNodes(
       PCollection<McfGraph> graph,
       SpannerClient spannerClient,

pipeline/ingestion/src/main/java/org/datacommons/ingestion/data/Node.java

@@ -1,7 +1,7 @@
 package org.datacommons.ingestion.data;
 
-import com.google.cloud.ByteArray;
 import java.io.Serializable;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import org.apache.beam.sdk.coders.DefaultCoder;
@@ -16,10 +16,13 @@ public class Node implements Serializable {
 
   private String subjectId;
   private String value;
-  private ByteArray bytes;
+  private byte[] bytes;
   private String name;
   private List<String> types;
 
+  @SuppressWarnings("unused")
+  private Node() {}
+
   // Private constructor to enforce use of Builder
   private Node(Builder builder) {
     this.subjectId = builder.subjectId;
@@ -41,7 +44,7 @@ public String getValue() {
     return value;
   }
 
-  public ByteArray getBytes() {
+  public byte[] getBytes() {
     return bytes;
   }
 
@@ -60,7 +63,7 @@ public boolean equals(Object o) {
     Node node = (Node) o;
     return Objects.equals(subjectId, node.subjectId)
         && Objects.equals(value, node.value)
-        && Objects.equals(bytes, node.bytes)
+        && Arrays.equals(bytes, node.bytes)
         && Objects.equals(name, node.name)
         && Objects.equals(types, node.types);
   }
@@ -74,13 +77,13 @@ public int hashCode() {
   public String toString() {
     return String.format(
         "Node{subjectId='%s', value='%s', bytes='%s', name='%s', types=%s}",
-        subjectId, value, bytes, name, types);
+        subjectId, value, Arrays.toString(bytes), name, types);
   }
 
   public static class Builder {
     private String subjectId = "";
     private String value = "";
-    private ByteArray bytes = null;
+    private byte[] bytes = new byte[0];
     private String name = "";
     private List<String> types = List.of();
 
@@ -96,7 +99,7 @@ public Builder value(String value) {
       return this;
     }
 
-    public Builder bytes(ByteArray bytes) {
+    public Builder bytes(byte[] bytes) {
       this.bytes = bytes;
       return this;
     }

pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/ImportGroupPipeline.java

@@ -11,13 +11,16 @@
 import org.apache.beam.sdk.metrics.Counter;
 import org.apache.beam.sdk.metrics.Metrics;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.Flatten;
 import org.apache.beam.sdk.transforms.Values;
 import org.apache.beam.sdk.transforms.Wait;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionList;
 import org.apache.beam.sdk.values.PCollectionTuple;
+import org.datacommons.ingestion.data.Edge;
 import org.datacommons.ingestion.data.GraphReader;
+import org.datacommons.ingestion.data.Node;
 import org.datacommons.ingestion.spanner.SpannerClient;
 import org.datacommons.pipeline.util.PipelineUtils;
 import org.datacommons.proto.Mcf.McfGraph;
@@ -95,9 +98,10 @@ public static void main(String[] args) {
           GraphReader.getProvenanceMcf(
               options.getStorageBucketId(), importName, latestVersion, pipeline);
 
-      PCollection<Mutation> deleteMutations =
-          GraphReader.getDeleteMutations(importName, provenance, pipeline, spannerClient);
-      deleteMutationList.add(deleteMutations);
+      //   PCollection<Mutation> deleteMutations =
+      //       GraphReader.getDeleteMutations(importName, provenance, pipeline, spannerClient);
+      //   deleteMutationList.add(deleteMutations);
+
       // Read schema mcf files and combine MCF nodes, and convert to spanner mutations (Node/Edge).
       String graphPath =
           latestVersion.replaceAll("/+$", "")
@@ -115,13 +119,23 @@ public static void main(String[] args) {
               .and(provenanceMcf)
               .apply("FlattenSchema", Flatten.pCollections());
 
-      PCollection<Mutation> edgeMutations =
-          GraphReader.graphToEdges(schemaMcf, provenance, spannerClient, edgeCounter)
-              .apply("ExtractEdgeMutations", Values.create());
+      PCollection<Edge> edges = GraphReader.mcfToEdges(schemaMcf, provenance, edgeCounter);
+
+      PCollection<Mutation> edgeMutations = GraphReader.edgeToMutations(edges, spannerClient);
+
+      PCollection<Node> newNodes =
+          GraphReader.mcfToNodes(schemaMcf, nodeCounter, nodeInvalidTypeCounter);
+
+      PCollection<Node> existingNodes = spannerClient.readExistingNodes(newNodes);
+
+      PCollection<Node> mergedNodes =
+          PCollectionList.of(newNodes)
+              .and(existingNodes)
+              .apply("FlattenNodes", Flatten.pCollections());
+
+      PCollection<Node> finalNodes = GraphReader.combineNodes(mergedNodes);
 
-      PCollection<Mutation> nodeMutations =
-          GraphReader.graphToNodes(schemaMcf, spannerClient, nodeCounter, nodeInvalidTypeCounter)
-              .apply("ExtractEdgeMutations", Values.create());
+      PCollection<Mutation> nodeMutations = GraphReader.nodeToMutations(finalNodes, spannerClient);
 
       nodeMutationList.add(nodeMutations);
       edgeMutationList.add(edgeMutations);
@@ -135,8 +149,9 @@ public static void main(String[] args) {
       obsMutationList.add(observationMutations);
     }
     PCollection<Mutation> deleteMutations =
-        PCollectionList.of(deleteMutationList)
-            .apply("FlattenDeleteMutations", Flatten.pCollections());
+        pipeline.apply(Create.empty(org.apache.beam.sdk.values.TypeDescriptor.of(Mutation.class)));
+    // PCollectionList.of(deleteMutationList)
+    //     .apply("FlattenDeleteMutations", Flatten.pCollections());
     SpannerWriteResult deleted =
         deleteMutations.apply("DeleteImportsFromSpanner", spannerClient.getWriteTransform());
     // Write the mutations to spanner.