From 3140eb7851dd5300ae13c963e835f3be824bf53f Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Tue, 6 Oct 2026 12:52:57 -0700 Subject: [PATCH] Prep changes for merging ValidatesRunner jobs --- .../apache/beam/sdk/CompositePipeline.java | 158 +++++++++++++ .../java/org/apache/beam/sdk/Pipeline.java | 2 +- .../beam/sdk/runners/TransformHierarchy.java | 28 ++- .../construction/PipelineTranslation.java | 19 +- .../beam/sdk/CompositePipelineTest.java | 221 ++++++++++++++++++ .../sdk/runners/TransformHierarchyTest.java | 40 ++++ 6 files changed, 456 insertions(+), 12 deletions(-) create mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/CompositePipeline.java create mode 100644 sdks/java/core/src/test/java/org/apache/beam/sdk/CompositePipelineTest.java diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/CompositePipeline.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/CompositePipeline.java new file mode 100644 index 000000000000..855a95c9d7b7 --- /dev/null +++ b/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; + +/** + * For internal use only; no backwards-compatibility guarantees. + * + *

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. + * + *

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. + * + *

Visitor contract: {@link #traverseTopologically} invokes {@link PipelineVisitor#enterPipeline} + * and {@link PipelineVisitor#leavePipeline} once per member, 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()}. + * + *

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 memberPipelines; + + /** + * Creates a composite over {@code memberPipelines}. + * + * @throws IllegalArgumentException if two members share a top-level transform name + */ + public CompositePipeline(PipelineOptions options, List 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 getMemberPipelines() { + return memberPipelines; + } + + @Override + public void traverseTopologically(PipelineVisitor visitor) { + for (Pipeline member : memberPipelines) { + member.traverseTopologically(visitor); + } + } + + @Override + public void replaceAll(List overrides) { + for (Pipeline member : memberPipelines) { + member.replaceAll(overrides); + } + } + + /** + * @throws UnsupportedOperationException always; a composite is a read-only view. + */ + @Override + public OutputT apply(PTransform root) { + throw new UnsupportedOperationException("Transforms cannot be applied to a CompositePipeline"); + } + + /** + * @throws UnsupportedOperationException always; a composite is a read-only view. + */ + @Override + public OutputT apply( + String name, PTransform root) { + throw new UnsupportedOperationException("Transforms cannot be applied to a CompositePipeline"); + } + + private static void checkTopLevelNamesDisjoint(List members) { + Map 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 topLevelTransformNames(Pipeline pipeline) { + List 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; + } +} diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/Pipeline.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/Pipeline.java index 184052520321..8500a937588d 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/Pipeline.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/Pipeline.java @@ -525,7 +525,7 @@ public static OutputT applyTran private final List> errorHandlers = new ArrayList<>(); - private Pipeline(TransformHierarchy transforms, PipelineOptions options) { + protected Pipeline(TransformHierarchy transforms, PipelineOptions options) { CoderTranslation.verifyModelCodersRegistered(); this.transforms = transforms; this.defaultOptions = options; diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java index 77b22d25d23d..b20812e74a51 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java +++ b/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 unexpandedInputs; private final Map, 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. + * + *

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; diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/PipelineTranslation.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/PipelineTranslation.java index 828b4f7ad83c..752346b28a66 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/PipelineTranslation.java +++ b/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> 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); } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/CompositePipelineTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/CompositePipelineTest.java new file mode 100644 index 000000000000..d7359148abe4 --- /dev/null +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/CompositePipelineTest.java @@ -0,0 +1,221 @@ +/* + * 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.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.contains; +import static org.hamcrest.Matchers.containsString; +import static org.hamcrest.Matchers.everyItem; +import static org.hamcrest.Matchers.hasItem; +import static org.hamcrest.Matchers.hasItems; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.not; +import static org.hamcrest.Matchers.startsWith; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.sdk.Pipeline.PipelineVisitor; +import org.apache.beam.sdk.io.GenerateSequence; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.runners.PTransformOverride; +import org.apache.beam.sdk.runners.TransformHierarchy; +import org.apache.beam.sdk.runners.TransformHierarchy.Node; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.resourcehints.ResourceHints; +import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.TypeDescriptors; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link CompositePipeline}. */ +@RunWith(JUnit4.class) +public class CompositePipelineTest { + + private static final PipelineOptions OPTIONS = PipelineOptionsFactory.create(); + + /** A pipeline whose root is named {@code rootName} (or unnamed if empty). */ + private static Pipeline pipelineWithRoot(String rootName) { + TransformHierarchy hierarchy = new TransformHierarchy(ResourceHints.create()); + if (!rootName.isEmpty()) { + hierarchy.setRootName(rootName); + } + return Pipeline.forTransformHierarchy(hierarchy, PipelineOptionsFactory.create()); + } + + /** Applies {@code Create -> Map} with fixed names and a {@link PAssert} to {@code pipeline}. */ + private static PCollection applyCreateAndMap(Pipeline pipeline, Integer... values) { + PCollection out = + pipeline + .apply("Create", Create.of(ImmutableList.copyOf(values))) + .apply("Map", MapElements.into(TypeDescriptors.integers()).via(x -> x + 1)); + PAssert.that(out).empty(); // never run; only counted + return out; + } + + /** Full names of every non-root transform reached by traversing {@code pipeline}. */ + private static List fullNames(Pipeline pipeline) { + List names = new ArrayList<>(); + pipeline.traverseTopologically( + new PipelineVisitor.Defaults() { + @Override + public CompositeBehavior enterCompositeTransform(Node node) { + if (!node.isRootNode()) { + names.add(node.getFullName()); + } + return CompositeBehavior.ENTER_TRANSFORM; + } + + @Override + public void visitPrimitiveTransform(Node node) { + names.add(node.getFullName()); + } + }); + return names; + } + + @Test + public void testTraversalCoversEveryMemberWithScopedNames() { + Pipeline a = pipelineWithRoot("a"); + Pipeline b = pipelineWithRoot("b"); + PCollection outA = applyCreateAndMap(a, 1, 2); + PCollection outB = applyCreateAndMap(b, 3); + + CompositePipeline composite = new CompositePipeline(OPTIONS, ImmutableList.of(a, b)); + assertThat(composite.getMemberPipelines(), contains(a, b)); + + List names = fullNames(composite); + assertThat(names, hasItems("a/Create", "a/Map", "b/Create", "b/Map")); + assertThat(names, not(hasItem("Create"))); + // One-shot visitors such as PAssert's counter are applied per member, not to the composite. + assertEquals(1, PAssert.countAsserts(a)); + assertEquals(1, PAssert.countAsserts(b)); + // PCollection names follow their producing transform's full name, so they are scoped too. + assertThat(outA.getName(), startsWith("a/")); + assertThat(outB.getName(), startsWith("b/")); + } + + @Test + public void testTranslatesToProtoWithDisjointUniqueNames() { + Pipeline a = pipelineWithRoot("a"); + Pipeline b = pipelineWithRoot("b"); + applyCreateAndMap(a, 1); + applyCreateAndMap(b, 2); + + RunnerApi.Pipeline proto = + PipelineTranslation.toProto(new CompositePipeline(OPTIONS, ImmutableList.of(a, b))); + + List uniqueNames = new ArrayList<>(); + for (RunnerApi.PTransform transform : proto.getComponents().getTransformsMap().values()) { + uniqueNames.add(transform.getUniqueName()); + } + assertThat(uniqueNames, hasItems("a/Create", "a/Map", "b/Create", "b/Map")); + // Both members' top-level transforms are roots of the merged graph. + List rootNames = new ArrayList<>(); + for (String id : proto.getRootTransformIdsList()) { + rootNames.add(proto.getComponents().getTransformsOrThrow(id).getUniqueName()); + } + assertThat(rootNames, hasItems("a/Create", "a/Map", "b/Create", "b/Map")); + assertThat(rootNames, everyItem(not(startsWith("/")))); + } + + @Test + public void testReplaceAllReachesEveryMember() { + Pipeline a = pipelineWithRoot("a"); + Pipeline b = pipelineWithRoot("b"); + a.apply("Seq", GenerateSequence.from(0).to(10)); + b.apply("Seq", GenerateSequence.from(0).to(10)); + CompositePipeline composite = new CompositePipeline(OPTIONS, ImmutableList.of(a, b)); + + composite.replaceAll( + ImmutableList.of( + PTransformOverride.of( + application -> application.getTransform() instanceof GenerateSequence, + new PipelineTest.GenerateSequenceToCreateOverride()))); + + Map> topLevel = new HashMap<>(); + composite.traverseTopologically( + new PipelineVisitor.Defaults() { + @Override + public CompositeBehavior enterCompositeTransform(Node node) { + if (!node.isRootNode() && node.getEnclosingNode().isRootNode()) { + topLevel.put(node.getFullName(), node.getTransform()); + } + return CompositeBehavior.ENTER_TRANSFORM; + } + }); + assertThat(topLevel.get("a/Seq"), instanceOf(Create.Values.class)); + assertThat(topLevel.get("b/Seq"), instanceOf(Create.Values.class)); + // Members were replaced in place, so a member on its own sees the replacement too. + assertThat(fullNames(a), everyItem(startsWith("a/"))); + } + + @Test + public void testRejectsMembersWithCollidingTopLevelNames() { + Pipeline unnamed0 = pipelineWithRoot(""); + Pipeline unnamed1 = pipelineWithRoot(""); + applyCreateAndMap(unnamed0, 1); + applyCreateAndMap(unnamed1, 2); + IllegalArgumentException e = + assertThrows( + IllegalArgumentException.class, + () -> new CompositePipeline(OPTIONS, ImmutableList.of(unnamed0, unnamed1))); + assertThat(e.getMessage(), containsString("\"Create\"")); + assertThat(e.getMessage(), containsString("setRootName")); + + Pipeline sameRoot0 = pipelineWithRoot("same"); + Pipeline sameRoot1 = pipelineWithRoot("same"); + applyCreateAndMap(sameRoot0, 1); + applyCreateAndMap(sameRoot1, 2); + assertThrows( + IllegalArgumentException.class, + () -> new CompositePipeline(OPTIONS, ImmutableList.of(sameRoot0, sameRoot1))); + } + + @Test + public void testAcceptsMembersWithDistinctNamesAndNoRoot() { + // Root names are the easy way to get disjoint names, but any disjoint naming is accepted. + Pipeline p0 = pipelineWithRoot(""); + Pipeline p1 = pipelineWithRoot(""); + p0.apply("First", Create.of(1)); + p1.apply("Second", Create.of(2)); + assertThat( + fullNames(new CompositePipeline(OPTIONS, ImmutableList.of(p0, p1))), + hasItems("First", "Second")); + } + + @Test + public void testApplyIsUnsupported() { + Pipeline a = pipelineWithRoot("a"); + applyCreateAndMap(a, 1); + CompositePipeline composite = new CompositePipeline(OPTIONS, ImmutableList.of(a)); + assertThrows(UnsupportedOperationException.class, () -> composite.apply(Create.of(1))); + assertThrows(UnsupportedOperationException.class, () -> composite.apply("Named", Create.of(1))); + } +} diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java index ad775643f904..1769c6da33f7 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java @@ -106,6 +106,46 @@ public void pushThenPopSucceeds() { assertThat(hierarchy.getCurrent(), equalTo(root)); } + @Test + public void rootIsUnnamedByDefault() { + assertThat(hierarchy.getCurrent().isRootNode(), is(true)); + assertThat(hierarchy.getCurrent().getFullName(), equalTo("")); + } + + @Test + public void setRootNameRenamesRoot() { + hierarchy.setRootName("myRoot"); + TransformHierarchy.Node root = hierarchy.getCurrent(); + assertThat(root.isRootNode(), is(true)); + // Pipeline builds every full name from getCurrent().getFullName(), so naming the root is what + // makes all subsequently applied transforms "myRoot/...". + assertThat(root.getFullName(), equalTo("myRoot")); + + TransformHierarchy.Node node = + hierarchy.pushNode("myRoot/Create", PBegin.in(pipeline), Create.of(1)); + assertThat(node.getEnclosingNode(), equalTo(root)); + assertThat(node.getFullName(), equalTo("myRoot/Create")); + hierarchy.popNode(); + assertThat(hierarchy.getCurrent(), equalTo(root)); + assertThat(hierarchy.getCurrent().getFullName(), equalTo("myRoot")); + } + + @Test + public void setRootNameAfterPushFails() { + hierarchy.pushNode("Create", PBegin.in(pipeline), Create.of(1)); + hierarchy.popNode(); + thrown.expect(IllegalStateException.class); + thrown.expectMessage("already contains transforms"); + hierarchy.setRootName("tooLate"); + } + + @Test + public void setRootNameWhileInsideCompositeFails() { + hierarchy.pushNode("Create", PBegin.in(pipeline), Create.of(1)); + thrown.expect(IllegalStateException.class); + hierarchy.setRootName("tooLate"); + } + @Test public void emptyCompositeSucceeds() { PCollection created =