Skip flattening for mutations across imports

Author: vish-cs

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

@@ -4,18 +4,20 @@
 import com.google.gson.JsonArray;
 import com.google.gson.JsonElement;
 import com.google.gson.JsonParser;
-import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.io.gcp.spanner.SpannerWriteResult;
 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.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionList;
 import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.TypeDescriptor;
 import org.datacommons.ingestion.spanner.SpannerClient;
 import org.datacommons.ingestion.util.GraphReader;
 import org.datacommons.ingestion.util.PipelineUtils;
@@ -70,12 +72,6 @@ public static void buildPipeline(
       Pipeline pipeline, IngestionPipelineOptions options, SpannerClient spannerClient) {
     LOGGER.info("Running import pipeline for imports: {}", options.getImportList());
 
-    // Initialize lists to hold mutations from all imports.
-    List<PCollection<Void>> deleteOpsList = new ArrayList<>();
-    List<PCollection<Mutation>> obsMutationList = new ArrayList<>();
-    List<PCollection<Mutation>> edgeMutationList = new ArrayList<>();
-    List<PCollection<Mutation>> nodeMutationList = new ArrayList<>();
-
     // Parse the input import list JSON.
     JsonElement jsonElement = JsonParser.parseString(options.getImportList());
     JsonArray jsonArray = jsonElement.getAsJsonArray();
@@ -97,37 +93,8 @@ public static void buildPipeline(
       String graphPath = pathElement.getAsString();
 
       // Process the individual import.
-      processImport(
-          pipeline,
-          spannerClient,
-          importName,
-          graphPath,
-          options.getSkipDelete(),
-          deleteOpsList,
-          nodeMutationList,
-          edgeMutationList,
-          obsMutationList);
-    }
-    // Finally, aggregate all collected mutations and write them to Spanner.
-    // 1. Process Deletes:
-    // First, execute all delete mutations to clear old data for the imports.
-    PCollection<Void> deleted =
-        PCollectionList.of(deleteOpsList).apply("DeleteOps", Flatten.pCollections());
-
-    // 2. Process Observations:
-    // Write observation mutations after deletes are complete.
-    if (options.getWriteObsGraph()) {
-      spannerClient.writeMutations(pipeline, "Observations", obsMutationList, deleted);
+      processImport(pipeline, spannerClient, importName, graphPath, options.getSkipDelete());
     }
-
-    // 3. Process Nodes:
-    // Write node mutations after deletes are complete.
-    SpannerWriteResult writtenNodes =
-        spannerClient.writeMutations(pipeline, "Nodes", nodeMutationList, deleted);
-
-    // 4. Process Edges:
-    // Write edge mutations only after node mutations are complete to ensure referential integrity.
-    spannerClient.writeMutations(pipeline, "Edges", edgeMutationList, writtenNodes.getOutput());
   }
 
   /**
@@ -138,31 +105,33 @@ public static void buildPipeline(
    * @param importName The name of the import.
    * @param graphPath The full path to the graph data.
    * @param skipDelete Whether to skip delete operations.
-   * @param deleteOpsList List to collect delete Ops.
-   * @param nodeMutationList List to collect node mutations.
-   * @param edgeMutationList List to collect edge mutations.
-   * @param obsMutationList List to collect observation mutations.
    */
   private static void processImport(
       Pipeline pipeline,
       SpannerClient spannerClient,
       String importName,
       String graphPath,
-      boolean skipDelete,
-      List<PCollection<Void>> deleteOpsList,
-      List<PCollection<Mutation>> nodeMutationList,
-      List<PCollection<Mutation>> edgeMutationList,
-      List<PCollection<Mutation>> obsMutationList) {
+      boolean skipDelete) {
     LOGGER.info("Import: {} Graph path: {}", importName, graphPath);
 
     String provenance = "dc/base/" + importName;
 
     // 1. Prepare Deletes:
     // Generate mutations to delete existing data for this import/provenance.
+    // Create a dummy signal if deletes are skipped, so downstream dependencies are satisfied
+    // immediately.
+    PCollection<Void> deleteObsWait;
+    PCollection<Void> deleteEdgesWait;
     if (!skipDelete) {
-      List<PCollection<Void>> deleteOps =
-          GraphReader.deleteExistingDataForImport(importName, provenance, pipeline, spannerClient);
-      deleteOpsList.addAll(deleteOps);
+      deleteObsWait = spannerClient.deleteObservationsForImport(importName, pipeline);
+      deleteEdgesWait = spannerClient.deleteEdgesForImport(provenance, pipeline);
+    } else {
+      deleteObsWait =
+          pipeline.apply(
+              "CreateEmptyObsWait-" + importName, Create.empty(TypeDescriptor.of(Void.class)));
+      deleteEdgesWait =
+          pipeline.apply(
+              "CreateEmptyEdgesWait-" + importName, Create.empty(TypeDescriptor.of(Void.class)));
     }
 
     // 2. Read and Split Graph:
@@ -176,29 +145,50 @@ private static void processImport(
     PCollection<McfGraph> schemaNodes = graphNodes.get(PipelineUtils.SCHEMA_NODES_TAG);
 
     // 3. Process Schema Nodes:
-    // Combine schema nodes if required, then convert to Node and Edge mutations.
+    // Combine nodes if required.
     PCollection<McfGraph> combinedGraph = schemaNodes;
     if (IMPORTS_TO_COMBINE.contains(importName)) {
       combinedGraph = PipelineUtils.combineGraphNodes(schemaNodes);
     }
+
+    // Convert all nodes to mutations
     PCollection<Mutation> nodeMutations =
         GraphReader.graphToNodes(
-                importName, combinedGraph, spannerClient, nodeCounter, nodeInvalidTypeCounter)
+                "NodeMutations-" + importName,
+                combinedGraph,
+                spannerClient,
+                nodeCounter,
+                nodeInvalidTypeCounter)
             .apply("ExtractNodeMutations-" + importName, Values.create());
     PCollection<Mutation> edgeMutations =
-        GraphReader.graphToEdges(importName, combinedGraph, provenance, spannerClient, edgeCounter)
+        GraphReader.graphToEdges(
+                "EdgeMutations-" + importName,
+                combinedGraph,
+                provenance,
+                spannerClient,
+                edgeCounter)
             .apply("ExtractEdgeMutations-" + importName, Values.create());
 
-    nodeMutationList.add(nodeMutations);
-    edgeMutationList.add(edgeMutations);
+    // Write Nodes
+    SpannerWriteResult writtenNodes =
+        spannerClient.writeMutations(pipeline, "Nodes-" + importName, List.of(nodeMutations), null);
+
+    PCollection<Void> writeEdgesWait =
+        PCollectionList.of(Arrays.asList(writtenNodes.getOutput(), deleteEdgesWait))
+            .apply("FlattenDeleteOps-" + importName, Flatten.pCollections());
+    // Write Edges (wait for Nodes and deletes)
+    spannerClient.writeMutations(
+        pipeline, "Edges-" + importName, List.of(edgeMutations), writeEdgesWait);
 
     // 4. Process Observation Nodes:
     // Build an optimized graph from observation nodes and convert to Observation mutations.
     PCollection<McfOptimizedGraph> optimizedGraph =
         PipelineUtils.buildOptimizedMcfGraph(observationNodes);
     PCollection<Mutation> observationMutations =
         GraphReader.graphToObservations(optimizedGraph, importName, spannerClient, obsCounter)
-            .apply("ExtractObsMutations", Values.create());
-    obsMutationList.add(observationMutations);
+            .apply("ExtractObsMutations-" + importName, Values.create());
+    // Write Observations (wait for delete)
+    spannerClient.writeMutations(
+        pipeline, "Observations-" + importName, List.of(observationMutations), deleteObsWait);
   }
 }

pipeline/ingestion/src/test/java/org/datacommons/ingestion/pipeline/GraphIngestionPipelineIntegrationTest.java

@@ -23,7 +23,6 @@
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
-import java.util.UUID;
 import org.apache.beam.runners.dataflow.DataflowRunner;
 import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
 import org.apache.beam.sdk.Pipeline;
@@ -86,7 +85,7 @@ public class GraphIngestionPipelineIntegrationTest {
   private String region;
   private String emulatorHost;
   private boolean isLocal;
-  private String importName = "TestImport-" + UUID.randomUUID().toString();
+  private String importName = "TestImport";
   private String nodeNameValue = "Test Node Name";
   private SpannerClient spannerClient;
 

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

@@ -5,11 +5,9 @@
 import java.io.Serializable;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.metrics.Counter;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.ParDo;
@@ -97,6 +95,9 @@ public static List<Edge> graphToEdges(McfGraph graph, String provenance) {
         String dcid = GraphUtils.getPropertyValue(pv, "dcid");
         String subjectId = !dcid.isEmpty() ? dcid : McfUtil.stripNamespace(nodeEntry.getKey());
         for (Map.Entry<String, McfGraph.Values> entry : pv.entrySet()) {
+          if (entry.getKey().equals("dcid")) {
+            continue;
+          }
           for (TypedValue val : entry.getValue().getTypedValuesList()) {
             if (val.getType() != ValueType.RESOLVED_REF) {
               int valSize = val.getValue().getBytes(StandardCharsets.UTF_8).length;
@@ -155,13 +156,6 @@ public static Observation graphToObservations(McfOptimizedGraph graph, String im
     return obs.build();
   }
 
-  public static List<PCollection<Void>> deleteExistingDataForImport(
-      String importName, String provenance, Pipeline pipeline, SpannerClient spannerClient) {
-    return Arrays.asList(
-        spannerClient.deleteObservationsForImport(importName, pipeline),
-        spannerClient.deleteEdgesForImport(provenance, pipeline));
-  }
-
   public static PCollection<KV<String, Mutation>> graphToObservations(
       PCollection<McfOptimizedGraph> graph,
       String importName,

pipeline/util/src/test/java/org/datacommons/ingestion/util/GraphReaderTest.java

@@ -173,6 +173,14 @@ public void testGraphToEdges() {
                                     .setType(ValueType.TEXT)
                                     .setValue("Subject Node"))
                             .build())
+                    .putPvs(
+                        "dcid",
+                        McfGraph.Values.newBuilder()
+                            .addTypedValues(
+                                TypedValue.newBuilder()
+                                    .setType(ValueType.TEXT)
+                                    .setValue("dcid_subject"))
+                            .build())
                     .putPvs(
                         "typeOf",
                         McfGraph.Values.newBuilder()