Prep changes for merging ValidatesRunner jobs

sdks/java/core/src/main/java/org/apache/beam/sdk/CompositePipeline.java

@@ -0,0 +1,158 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk;
+
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.runners.PTransformOverride;
+import org.apache.beam.sdk.runners.TransformHierarchy;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.POutput;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+
+/**
+ * <b><i>For internal use only; no backwards-compatibility guarantees.</i></b>
+ *
+ * <p>A read-only {@link Pipeline} that presents several independently constructed member pipelines
+ * as one graph, so that a runner can translate and execute them together as a single job.
+ *
+ * <p>Only the two operations runners use to translate a pipeline are composed: {@link
+ * #traverseTopologically} visits every member in turn, and {@link #replaceAll} applies the runner's
+ * overrides to every member. Everything else is the composite's own: {@link #getOptions()} are the
+ * options passed to the constructor (member options are ignored), the coder and schema registries
+ * are fresh and empty, and transforms cannot be applied to it.
+ *
+ * <p>Visitor contract: {@link #traverseTopologically} invokes {@link PipelineVisitor#enterPipeline}
+ * and {@link PipelineVisitor#leavePipeline} once <em>per member</em>, each time with that member,
+ * and visits the member's root node in between; it never passes the composite itself to the
+ * visitor. Consequently {@link PipelineVisitor.Defaults#getPipeline()} always yields the member
+ * that owns the node being visited, which is what visitors needing a member's state (for example
+ * its {@link #getCoderRegistry()}) should use. Visitors that assume a single {@code
+ * enterPipeline}/root-node visit per traversal must instead be applied to each member in {@link
+ * #getMemberPipelines()}.
+ *
+ * <p>Members must have pairwise distinct top-level transform names (and therefore distinct full
+ * names throughout, since every nested name and {@link org.apache.beam.sdk.values.PCollection} name
+ * is derived from its enclosing transform's); this is checked at construction. The simplest way to
+ * guarantee it is to give each member a unique root name via {@link TransformHierarchy#setRootName}
+ * before applying any transform.
+ */
+@Internal
+public final class CompositePipeline extends Pipeline {
+
+  private final List<Pipeline> memberPipelines;
+
+  /**
+   * Creates a composite over {@code memberPipelines}.
+   *
+   * @throws IllegalArgumentException if two members share a top-level transform name
+   */
+  public CompositePipeline(PipelineOptions options, List<Pipeline> memberPipelines) {
+    super(options);
+    this.memberPipelines = ImmutableList.copyOf(memberPipelines);
+    checkTopLevelNamesDisjoint(this.memberPipelines);
+  }
+
+  /** The member pipelines, in the order they were given and are traversed. */
+  public List<Pipeline> getMemberPipelines() {
+    return memberPipelines;
+  }
+
+  @Override
+  public void traverseTopologically(PipelineVisitor visitor) {
+    for (Pipeline member : memberPipelines) {
+      member.traverseTopologically(visitor);
+    }
+  }
+
+  @Override
+  public void replaceAll(List<PTransformOverride> overrides) {
+    for (Pipeline member : memberPipelines) {
+      member.replaceAll(overrides);
+    }
+  }
+
+  /**
+   * @throws UnsupportedOperationException always; a composite is a read-only view.
+   */
+  @Override
+  public <OutputT extends POutput> OutputT apply(PTransform<? super PBegin, OutputT> root) {
+    throw new UnsupportedOperationException("Transforms cannot be applied to a CompositePipeline");
+  }
+
+  /**
+   * @throws UnsupportedOperationException always; a composite is a read-only view.
+   */
+  @Override
+  public <OutputT extends POutput> OutputT apply(
+      String name, PTransform<? super PBegin, OutputT> root) {
+    throw new UnsupportedOperationException("Transforms cannot be applied to a CompositePipeline");
+  }
+
+  private static void checkTopLevelNamesDisjoint(List<Pipeline> members) {
+    Map<String, Integer> owners = new HashMap<>();
+    for (int i = 0; i < members.size(); i++) {
+      for (String name : topLevelTransformNames(members.get(i))) {
+        Integer previous = owners.putIfAbsent(name, i);
+        checkArgument(
+            previous == null,
+            "Member pipelines %s and %s both have a top-level transform named \"%s\"; give each"
+                + " member a unique root name (see TransformHierarchy#setRootName) so that their"
+                + " transform names are disjoint",
+            previous,
+            i,
+            name);
+      }
+    }
+  }
+
+  private static List<String> topLevelTransformNames(Pipeline pipeline) {
+    List<String> names = new ArrayList<>();
+    pipeline.traverseTopologically(
+        new PipelineVisitor.Defaults() {
+          @Override
+          public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
+            if (isTopLevel(node)) {
+              names.add(node.getFullName());
+              return CompositeBehavior.DO_NOT_ENTER_TRANSFORM;
+            }
+            return CompositeBehavior.ENTER_TRANSFORM;
+          }
+
+          @Override
+          public void visitPrimitiveTransform(TransformHierarchy.Node node) {
+            if (isTopLevel(node)) {
+              names.add(node.getFullName());
+            }
+          }
+
+          private boolean isTopLevel(TransformHierarchy.Node node) {
+            TransformHierarchy.Node enclosing = node.getEnclosingNode();
+            return enclosing != null && enclosing.isRootNode();
+          }
+        });
+    return names;
+  }
+}

sdks/java/core/src/main/java/org/apache/beam/sdk/Pipeline.java

@@ -525,7 +525,7 @@ public static <InputT extends PInput, OutputT extends POutput> OutputT applyTran
 
   private final List<ErrorHandler<?, ?>> errorHandlers = new ArrayList<>();
 
-  private Pipeline(TransformHierarchy transforms, PipelineOptions options) {
+  protected Pipeline(TransformHierarchy transforms, PipelineOptions options) {
     CoderTranslation.verifyModelCodersRegistered();
     this.transforms = transforms;
     this.defaultOptions = options;

sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java

@@ -62,7 +62,7 @@
 public class TransformHierarchy {
   private static final Logger LOG = LoggerFactory.getLogger(TransformHierarchy.class);
 
-  private final Node root;
+  private Node root;
   private final Map<Node, PInput> unexpandedInputs;
   private final Map<PCollection<?>, Node> producers;
 
@@ -76,7 +76,24 @@ public TransformHierarchy(ResourceHints resourceHints) {
     producers = new HashMap<>();
     producerInput = new HashMap<>();
     unexpandedInputs = new HashMap<>();
-    root = new Node(resourceHints);
+    root = new Node(resourceHints, "");
+    current = root;
+  }
+
+  /**
+   * Names the root of this hierarchy. Callers that derive full transform names from {@code
+   * getCurrent().getFullName()}, as {@link org.apache.beam.sdk.Pipeline} does, thereby prefix the
+   * full name of every transform (and of every {@link PCollection} named after one) subsequently
+   * added with {@code rootName + "/"}, which lets the names of independently constructed
+   * hierarchies be kept disjoint.
+   *
+   * <p>May only be called while the hierarchy is still empty.
+   */
+  public void setRootName(String rootName) {
+    checkState(
+        root.parts.isEmpty() && current == root && unexpandedInputs.isEmpty(),
+        "Cannot name the root of a TransformHierarchy that already contains transforms");
+    root = new Node(root.resourceHints, rootName);
     current = root;
   }
 
@@ -264,12 +281,13 @@ public class Node {
 
     /**
      * Creates the root-level node. The root level node has a null enclosing node, a null transform,
-     * an empty map of inputs, an empty map of outputs, and a name equal to the empty string.
+     * an empty map of inputs, an empty map of outputs, and the given name (the empty string unless
+     * {@link TransformHierarchy#setRootName} was used).
      */
-    private Node(ResourceHints resourceHints) {
+    private Node(ResourceHints resourceHints, String fullName) {
       this.enclosingNode = null;
       this.transform = null;
-      this.fullName = "";
+      this.fullName = fullName;
       this.inputs = Collections.emptyMap();
       this.outputs = Collections.emptyMap();
       this.resourceHints = resourceHints;

sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/PipelineTranslation.java

@@ -70,6 +70,13 @@ public static RunnerApi.Pipeline toProto(
           private final ListMultimap<Node, AppliedPTransform<?, ?, ?>> children =
               ArrayListMultimap.create();
 
+          // Nodes are turned into AppliedPTransforms against getPipeline(), i.e. the pipeline most
+          // recently passed to enterPipeline, rather than the {@code pipeline} being translated.
+          // For an ordinary Pipeline these are the same object. For a
+          // {@link org.apache.beam.sdk.CompositePipeline}, which traverses each member pipeline in
+          // turn, getPipeline() is the member that owns the node, so that translators which consult
+          // the owning pipeline (e.g. CombineTranslation's use of the coder registry) see the
+          // member's state rather than the composite's empty registries.
           @Override
           public void leaveCompositeTransform(Node node) {
             if (node.isRootNode()) {
@@ -78,10 +85,10 @@ public void leaveCompositeTransform(Node node) {
               }
             } else {
               // TODO: Include DisplayData in the proto
-              children.put(node.getEnclosingNode(), node.toAppliedPTransform(pipeline));
+              AppliedPTransform<?, ?, ?> appliedTransform = node.toAppliedPTransform(getPipeline());
+              children.put(node.getEnclosingNode(), appliedTransform);
               try {
-                components.registerPTransform(
-                    node.toAppliedPTransform(pipeline), children.get(node));
+                components.registerPTransform(appliedTransform, children.get(node));
               } catch (IOException e) {
                 throw new RuntimeException(e);
               }
@@ -91,10 +98,10 @@ public void leaveCompositeTransform(Node node) {
           @Override
           public void visitPrimitiveTransform(Node node) {
             // TODO: Include DisplayData in the proto
-            children.put(node.getEnclosingNode(), node.toAppliedPTransform(pipeline));
+            AppliedPTransform<?, ?, ?> appliedTransform = node.toAppliedPTransform(getPipeline());
+            children.put(node.getEnclosingNode(), appliedTransform);
             try {
-              components.registerPTransform(
-                  node.toAppliedPTransform(pipeline), Collections.emptyList());
+              components.registerPTransform(appliedTransform, Collections.emptyList());
             } catch (IOException e) {
               throw new IllegalStateException(e);
             }