runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/QueryablePipeline.java - beam - Git at Google

 /*
  * 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.runners.core.construction.graph;

 import static org.apache.beam.runners.core.construction.PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.COMBINE_PER_KEY_EXTRACT_OUTPUTS_TRANSFORM_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.COMBINE_PER_KEY_MERGE_ACCUMULATORS_TRANSFORM_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.COMBINE_PER_KEY_PRECOMBINE_TRANSFORM_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.CREATE_VIEW_TRANSFORM_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.FLATTEN_TRANSFORM_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.IMPULSE_TRANSFORM_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.MAP_WINDOWS_TRANSFORM_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.PAR_DO_TRANSFORM_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.READ_TRANSFORM_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_PROCESS_ELEMENTS_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_PROCESS_KEYED_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.TEST_STREAM_TRANSFORM_URN;
 import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;

 import java.util.ArrayDeque;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.Deque;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
 import org.apache.beam.model.pipeline.v1.RunnerApi;
 import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
 import org.apache.beam.model.pipeline.v1.RunnerApi.Environment;
 import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
 import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
 import org.apache.beam.model.pipeline.v1.RunnerApi.ParDoPayload;
 import org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline;
 import org.apache.beam.runners.core.construction.Environments;
 import org.apache.beam.runners.core.construction.NativeTransforms;
 import org.apache.beam.runners.core.construction.PTransformTranslation;
 import org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode;
 import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode;
 import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.MutableNetwork;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.Network;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.NetworkBuilder;

 /**
  * A {@link Pipeline} which has additional methods to relate nodes in the graph relative to each
  * other.
  */
 public class QueryablePipeline {
   // TODO: Is it better to have the signatures here require nodes in almost all contexts, or should
   // they all take strings? Nodes gives some degree of type signalling that names might not, but
   // it's more painful to construct the node. However, right now the traversal is done starting
   // at the roots and using nodes everywhere based on what any query has returned.
   /**
    * Create a new {@link QueryablePipeline} based on the provided components.
    *
    * <p>The returned {@link QueryablePipeline} will contain only the primitive transforms present
    * within the provided components.
    */
   public static QueryablePipeline forPrimitivesIn(Components components) {
     return new QueryablePipeline(getPrimitiveTransformIds(components), components);
   }

   /**
    * Create a new {@link QueryablePipeline} which uses the root transform IDs and components of the
    * provided {@link Pipeline}.
    */
   public static QueryablePipeline forPipeline(RunnerApi.Pipeline p) {
     return forTransforms(p.getRootTransformIdsList(), p.getComponents());
   }

   /**
    * Create a new {@link QueryablePipeline} based on the provided components containing only the
    * provided {@code transformIds}.
    */
   public static QueryablePipeline forTransforms(
       Collection<String> transformIds, Components components) {
     return new QueryablePipeline(transformIds, components);
   }

   private final Components components;

   /**
    * The {@link Pipeline} represented by a {@link Network}.
    *
    * <p>This is a directed bipartite graph consisting of {@link PTransformNode PTransformNodes} and
    * {@link PCollectionNode PCollectionNodes}. Each {@link PCollectionNode} has exactly one in edge,
    * and an arbitrary number of out edges. Each {@link PTransformNode} has an arbitrary number of in
    * and out edges.
    *
    * <p>Parallel edges are permitted, as a {@link PCollectionNode} can be consumed by a single
    * {@link PTransformNode} any number of times with different local names.
    */
   private final Network<PipelineNode, PipelineEdge> pipelineNetwork;

   private QueryablePipeline(Collection<String> transformIds, Components components) {
     this.components = components;
     this.pipelineNetwork = buildNetwork(transformIds, this.components);
   }

   /** Produces a {@link RunnerApi.Components} which contains only primitive transforms. */
   @VisibleForTesting
   static Collection<String> getPrimitiveTransformIds(RunnerApi.Components components) {
     Collection<String> ids = new LinkedHashSet<>();

     for (Map.Entry<String, PTransform> transformEntry : components.getTransformsMap().entrySet()) {
       PTransform transform = transformEntry.getValue();
       boolean isPrimitive = isPrimitiveTransform(transform);
       if (isPrimitive) {
         // Sometimes "primitive" transforms have sub-transforms (and even deeper-nested
         // descendents), due to runners
         // either rewriting them in terms of runner-specific transforms, or SDKs constructing them
         // in terms of other
         // underlying transforms (see https://issues.apache.org/jira/browse/BEAM-5441).
         // We consider any "leaf" descendents of these "primitive" transforms to be the true
         // "primitives" that we
         // preserve here; in the common case, this is just the "primitive" itself, which has no
         // descendents).
         Deque<String> transforms = new ArrayDeque<>();
         transforms.push(transformEntry.getKey());
         while (!transforms.isEmpty()) {
           String id = transforms.pop();
           PTransform next = components.getTransformsMap().get(id);
           List<String> subtransforms = next.getSubtransformsList();
           if (subtransforms.isEmpty()) {
             ids.add(id);
           } else {
             transforms.addAll(subtransforms);
           }
         }
       }
     }
     return ids;
   }

   private static final Set<String> PRIMITIVE_URNS =
       ImmutableSet.of(
           PAR_DO_TRANSFORM_URN,
           FLATTEN_TRANSFORM_URN,
           GROUP_BY_KEY_TRANSFORM_URN,
           IMPULSE_TRANSFORM_URN,
           ASSIGN_WINDOWS_TRANSFORM_URN,
           TEST_STREAM_TRANSFORM_URN,
           MAP_WINDOWS_TRANSFORM_URN,
           READ_TRANSFORM_URN,
           CREATE_VIEW_TRANSFORM_URN,
           COMBINE_PER_KEY_PRECOMBINE_TRANSFORM_URN,
           COMBINE_PER_KEY_MERGE_ACCUMULATORS_TRANSFORM_URN,
           COMBINE_PER_KEY_EXTRACT_OUTPUTS_TRANSFORM_URN,
           SPLITTABLE_PROCESS_KEYED_URN,
           SPLITTABLE_PROCESS_ELEMENTS_URN,
           SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN,
           SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN);

   /** Returns true if the provided transform is a primitive. */
   private static boolean isPrimitiveTransform(PTransform transform) {
     String urn = PTransformTranslation.urnForTransformOrNull(transform);
     return PRIMITIVE_URNS.contains(urn) || NativeTransforms.isNative(transform);
   }

   private MutableNetwork<PipelineNode, PipelineEdge> buildNetwork(
       Collection<String> transformIds, Components components) {
     MutableNetwork<PipelineNode, PipelineEdge> network =
         NetworkBuilder.directed().allowsParallelEdges(true).allowsSelfLoops(false).build();
     Set<PCollectionNode> unproducedCollections = new HashSet<>();
     for (String transformId : transformIds) {
       PTransform transform = components.getTransformsOrThrow(transformId);
       PTransformNode transformNode =
           PipelineNode.pTransform(transformId, this.components.getTransformsOrThrow(transformId));
       network.addNode(transformNode);
       for (String produced : transform.getOutputsMap().values()) {
         PCollectionNode producedNode =
             PipelineNode.pCollection(produced, components.getPcollectionsOrThrow(produced));
         network.addNode(producedNode);
         network.addEdge(transformNode, producedNode, new PerElementEdge());
         checkArgument(
             network.inDegree(producedNode) == 1,
             "A %s should have exactly one producing %s, but found %s:\nPCollection:\n%s\nProducers:\n%s",
             PCollectionNode.class.getSimpleName(),
             PTransformNode.class.getSimpleName(),
             network.predecessors(producedNode).size(),
             producedNode,
             network.predecessors(producedNode));
         unproducedCollections.remove(producedNode);
       }
       for (Map.Entry<String, String> consumed : transform.getInputsMap().entrySet()) {
         // This loop may add an edge between the consumed PCollection and the current PTransform.
         // The local name of the transform must be used to determine the type of edge.
         String pcollectionId = consumed.getValue();
         PCollectionNode consumedNode =
             PipelineNode.pCollection(
                 pcollectionId, this.components.getPcollectionsOrThrow(pcollectionId));
         if (network.addNode(consumedNode)) {
           // This node has been added to the network for the first time, so it has no producer.
           unproducedCollections.add(consumedNode);
         }
         if (getLocalSideInputNames(transform).contains(consumed.getKey())) {
           network.addEdge(consumedNode, transformNode, new SingletonEdge());
         } else {
           network.addEdge(consumedNode, transformNode, new PerElementEdge());
         }
       }
     }
     checkArgument(
         unproducedCollections.isEmpty(),
         "%ss %s were consumed but never produced",
         PCollectionNode.class.getSimpleName(),
         unproducedCollections);
     return network;
   }

   public Collection<PTransformNode> getTransforms() {
     return pipelineNetwork.nodes().stream()
         .filter(PTransformNode.class::isInstance)
         .map(PTransformNode.class::cast)
         .collect(Collectors.toList());
   }

   public Iterable<PTransformNode> getTopologicallyOrderedTransforms() {
     return StreamSupport.stream(
             Networks.topologicalOrder(pipelineNetwork, Comparator.comparing(PipelineNode::getId))
                 .spliterator(),
             false)
         .filter(PTransformNode.class::isInstance)
         .map(PTransformNode.class::cast)
         .collect(Collectors.toList());
   }

   /**
    * Get the transforms that are roots of this {@link QueryablePipeline}. These are all nodes which
    * have no input {@link PCollection}.
    */
   public Set<PTransformNode> getRootTransforms() {
     return pipelineNetwork.nodes().stream()
         .filter(pipelineNode -> pipelineNetwork.inEdges(pipelineNode).isEmpty())
         .map(pipelineNode -> (PTransformNode) pipelineNode)
         .collect(Collectors.toSet());
   }

   public PTransformNode getProducer(PCollectionNode pcollection) {
     return (PTransformNode) Iterables.getOnlyElement(pipelineNetwork.predecessors(pcollection));
   }

   /**
    * Get all of the {@link PTransformNode PTransforms} which consume the provided {@link
    * PCollectionNode} on a per-element basis.
    *
    * <p>If a {@link PTransformNode} consumes a {@link PCollectionNode} on a per-element basis one or
    * more times, it will appear a single time in the result.
    *
    * <p>In theory, a transform may consume a single {@link PCollectionNode} in both a per-element
    * and singleton manner. If this is the case, the transform node is included in the result, as it
    * does consume the {@link PCollectionNode} on a per-element basis.
    */
   public Set<PTransformNode> getPerElementConsumers(PCollectionNode pCollection) {
     return pipelineNetwork.successors(pCollection).stream()
         .filter(
             consumer ->
                 pipelineNetwork.edgesConnecting(pCollection, consumer).stream()
                     .anyMatch(PipelineEdge::isPerElement))
         .map(pipelineNode -> (PTransformNode) pipelineNode)
         .collect(Collectors.toSet());
   }

   /**
    * Same as {@link #getPerElementConsumers(PCollectionNode)}, but returns transforms that consume
    * the collection as a singleton.
    */
   public Set<PTransformNode> getSingletonConsumers(PCollectionNode pCollection) {
     return pipelineNetwork.successors(pCollection).stream()
         .filter(
             consumer ->
                 pipelineNetwork.edgesConnecting(pCollection, consumer).stream()
                     .anyMatch(edge -> !edge.isPerElement()))
         .map(pipelineNode -> (PTransformNode) pipelineNode)
         .collect(Collectors.toSet());
   }

   /**
    * Gets each {@link PCollectionNode} that the provided {@link PTransformNode} consumes on a
    * per-element basis.
    */
   public Set<PCollectionNode> getPerElementInputPCollections(PTransformNode ptransform) {
     return pipelineNetwork.inEdges(ptransform).stream()
         .filter(PipelineEdge::isPerElement)
         .map(edge -> (PCollectionNode) pipelineNetwork.incidentNodes(edge).source())
         .collect(Collectors.toSet());
   }

   public Set<PCollectionNode> getOutputPCollections(PTransformNode ptransform) {
     return pipelineNetwork.successors(ptransform).stream()
         .map(pipelineNode -> (PCollectionNode) pipelineNode)
         .collect(Collectors.toSet());
   }

   public Components getComponents() {
     return components;
   }

   /**
    * Returns the {@link SideInputReference SideInputReferences} that the provided transform consumes
    * as side inputs.
    */
   public Collection<SideInputReference> getSideInputs(PTransformNode transform) {
     return getLocalSideInputNames(transform.getTransform()).stream()
         .map(
             localName -> {
               String transformId = transform.getId();
               PTransform transformProto = components.getTransformsOrThrow(transformId);
               String collectionId = transform.getTransform().getInputsOrThrow(localName);
               PCollection collection = components.getPcollectionsOrThrow(collectionId);
               return SideInputReference.of(
                   PipelineNode.pTransform(transformId, transformProto),
                   localName,
                   PipelineNode.pCollection(collectionId, collection));
             })
         .collect(Collectors.toSet());
   }

   public Collection<UserStateReference> getUserStates(PTransformNode transform) {
     return getLocalUserStateNames(transform.getTransform()).stream()
         .map(
             localName -> {
               String transformId = transform.getId();
               PTransform transformProto = components.getTransformsOrThrow(transformId);
               // Get the main input PCollection id.
               String collectionId =
                   transform
                       .getTransform()
                       .getInputsOrThrow(
                           Iterables.getOnlyElement(
                               Sets.difference(
                                   transform.getTransform().getInputsMap().keySet(),
                                   ImmutableSet.builder()
                                       .addAll(getLocalSideInputNames(transformProto))
                                       .addAll(getLocalTimerNames(transformProto))
                                       .build())));
               PCollection collection = components.getPcollectionsOrThrow(collectionId);
               return UserStateReference.of(
                   PipelineNode.pTransform(transformId, transformProto),
                   localName,
                   PipelineNode.pCollection(collectionId, collection));
             })
         .collect(Collectors.toSet());
   }

   public Collection<TimerReference> getTimers(PTransformNode transform) {
     return getLocalTimerNames(transform.getTransform()).stream()
         .map(
             localName -> {
               String transformId = transform.getId();
               PTransform transformProto = components.getTransformsOrThrow(transformId);
               return TimerReference.of(
                   PipelineNode.pTransform(transformId, transformProto), localName);
             })
         .collect(Collectors.toSet());
   }

   private Set<String> getLocalSideInputNames(PTransform transform) {
     if (PAR_DO_TRANSFORM_URN.equals(transform.getSpec().getUrn())) {
       try {
         return ParDoPayload.parseFrom(transform.getSpec().getPayload()).getSideInputsMap().keySet();
       } catch (InvalidProtocolBufferException e) {
         throw new RuntimeException(e);
       }
     } else {
       return Collections.emptySet();
     }
   }

   private Set<String> getLocalUserStateNames(PTransform transform) {
     if (PAR_DO_TRANSFORM_URN.equals(transform.getSpec().getUrn())) {
       try {
         return ParDoPayload.parseFrom(transform.getSpec().getPayload()).getStateSpecsMap().keySet();
       } catch (InvalidProtocolBufferException e) {
         throw new RuntimeException(e);
       }
     } else {
       return Collections.emptySet();
     }
   }

   private Set<String> getLocalTimerNames(PTransform transform) {
     if (PAR_DO_TRANSFORM_URN.equals(transform.getSpec().getUrn())) {
       try {
         return ParDoPayload.parseFrom(transform.getSpec().getPayload()).getTimerSpecsMap().keySet();
       } catch (InvalidProtocolBufferException e) {
         throw new RuntimeException(e);
       }
     } else {
       return Collections.emptySet();
     }
   }

   public Optional<Environment> getEnvironment(PTransformNode parDo) {
     return Environments.getEnvironment(parDo.getId(), components);
   }

   private interface PipelineEdge {
     boolean isPerElement();
   }

   private static class PerElementEdge implements PipelineEdge {
     @Override
     public boolean isPerElement() {
       return true;
     }
   }

   private static class SingletonEdge implements PipelineEdge {
     @Override
     public boolean isPerElement() {
       return false;
     }
   }
 }
	/*
	* 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.runners.core.construction.graph;

	import static org.apache.beam.runners.core.construction.PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.COMBINE_PER_KEY_EXTRACT_OUTPUTS_TRANSFORM_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.COMBINE_PER_KEY_MERGE_ACCUMULATORS_TRANSFORM_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.COMBINE_PER_KEY_PRECOMBINE_TRANSFORM_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.CREATE_VIEW_TRANSFORM_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.FLATTEN_TRANSFORM_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.IMPULSE_TRANSFORM_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.MAP_WINDOWS_TRANSFORM_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.PAR_DO_TRANSFORM_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.READ_TRANSFORM_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_PROCESS_ELEMENTS_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_PROCESS_KEYED_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN;
	import static org.apache.beam.runners.core.construction.PTransformTranslation.TEST_STREAM_TRANSFORM_URN;
	import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;

	import java.util.ArrayDeque;
	import java.util.Collection;
	import java.util.Collections;
	import java.util.Comparator;
	import java.util.Deque;
	import java.util.HashSet;
	import java.util.LinkedHashSet;
	import java.util.List;
	import java.util.Map;
	import java.util.Optional;
	import java.util.Set;
	import java.util.stream.Collectors;
	import java.util.stream.StreamSupport;
	import org.apache.beam.model.pipeline.v1.RunnerApi;
	import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
	import org.apache.beam.model.pipeline.v1.RunnerApi.Environment;
	import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
	import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
	import org.apache.beam.model.pipeline.v1.RunnerApi.ParDoPayload;
	import org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline;
	import org.apache.beam.runners.core.construction.Environments;
	import org.apache.beam.runners.core.construction.NativeTransforms;
	import org.apache.beam.runners.core.construction.PTransformTranslation;
	import org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode;
	import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode;
	import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.MutableNetwork;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.Network;
	import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.NetworkBuilder;

	/**
	* A {@link Pipeline} which has additional methods to relate nodes in the graph relative to each
	* other.
	*/
	public class QueryablePipeline {
	// TODO: Is it better to have the signatures here require nodes in almost all contexts, or should
	// they all take strings? Nodes gives some degree of type signalling that names might not, but
	// it's more painful to construct the node. However, right now the traversal is done starting
	// at the roots and using nodes everywhere based on what any query has returned.
	/**
	* Create a new {@link QueryablePipeline} based on the provided components.
	*
	* <p>The returned {@link QueryablePipeline} will contain only the primitive transforms present
	* within the provided components.
	*/
	public static QueryablePipeline forPrimitivesIn(Components components) {
	return new QueryablePipeline(getPrimitiveTransformIds(components), components);
	}

	/**
	* Create a new {@link QueryablePipeline} which uses the root transform IDs and components of the
	* provided {@link Pipeline}.
	*/
	public static QueryablePipeline forPipeline(RunnerApi.Pipeline p) {
	return forTransforms(p.getRootTransformIdsList(), p.getComponents());
	}

	/**
	* Create a new {@link QueryablePipeline} based on the provided components containing only the
	* provided {@code transformIds}.
	*/
	public static QueryablePipeline forTransforms(
	Collection<String> transformIds, Components components) {
	return new QueryablePipeline(transformIds, components);
	}

	private final Components components;

	/**
	* The {@link Pipeline} represented by a {@link Network}.
	*
	* <p>This is a directed bipartite graph consisting of {@link PTransformNode PTransformNodes} and
	* {@link PCollectionNode PCollectionNodes}. Each {@link PCollectionNode} has exactly one in edge,
	* and an arbitrary number of out edges. Each {@link PTransformNode} has an arbitrary number of in
	* and out edges.
	*
	* <p>Parallel edges are permitted, as a {@link PCollectionNode} can be consumed by a single
	* {@link PTransformNode} any number of times with different local names.
	*/
	private final Network<PipelineNode, PipelineEdge> pipelineNetwork;

	private QueryablePipeline(Collection<String> transformIds, Components components) {
	this.components = components;
	this.pipelineNetwork = buildNetwork(transformIds, this.components);
	}

	/** Produces a {@link RunnerApi.Components} which contains only primitive transforms. */
	@VisibleForTesting
	static Collection<String> getPrimitiveTransformIds(RunnerApi.Components components) {
	Collection<String> ids = new LinkedHashSet<>();

	for (Map.Entry<String, PTransform> transformEntry : components.getTransformsMap().entrySet()) {
	PTransform transform = transformEntry.getValue();
	boolean isPrimitive = isPrimitiveTransform(transform);
	if (isPrimitive) {
	// Sometimes "primitive" transforms have sub-transforms (and even deeper-nested
	// descendents), due to runners
	// either rewriting them in terms of runner-specific transforms, or SDKs constructing them
	// in terms of other
	// underlying transforms (see https://issues.apache.org/jira/browse/BEAM-5441).
	// We consider any "leaf" descendents of these "primitive" transforms to be the true
	// "primitives" that we
	// preserve here; in the common case, this is just the "primitive" itself, which has no
	// descendents).
	Deque<String> transforms = new ArrayDeque<>();
	transforms.push(transformEntry.getKey());
	while (!transforms.isEmpty()) {
	String id = transforms.pop();
	PTransform next = components.getTransformsMap().get(id);
	List<String> subtransforms = next.getSubtransformsList();
	if (subtransforms.isEmpty()) {
	ids.add(id);
	} else {
	transforms.addAll(subtransforms);
	}
	}
	}
	}
	return ids;
	}

	private static final Set<String> PRIMITIVE_URNS =
	ImmutableSet.of(
	PAR_DO_TRANSFORM_URN,
	FLATTEN_TRANSFORM_URN,
	GROUP_BY_KEY_TRANSFORM_URN,
	IMPULSE_TRANSFORM_URN,
	ASSIGN_WINDOWS_TRANSFORM_URN,
	TEST_STREAM_TRANSFORM_URN,
	MAP_WINDOWS_TRANSFORM_URN,
	READ_TRANSFORM_URN,
	CREATE_VIEW_TRANSFORM_URN,
	COMBINE_PER_KEY_PRECOMBINE_TRANSFORM_URN,
	COMBINE_PER_KEY_MERGE_ACCUMULATORS_TRANSFORM_URN,
	COMBINE_PER_KEY_EXTRACT_OUTPUTS_TRANSFORM_URN,
	SPLITTABLE_PROCESS_KEYED_URN,
	SPLITTABLE_PROCESS_ELEMENTS_URN,
	SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN,
	SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN);

	/** Returns true if the provided transform is a primitive. */
	private static boolean isPrimitiveTransform(PTransform transform) {
	String urn = PTransformTranslation.urnForTransformOrNull(transform);
	return PRIMITIVE_URNS.contains(urn) \|\| NativeTransforms.isNative(transform);
	}

	private MutableNetwork<PipelineNode, PipelineEdge> buildNetwork(
	Collection<String> transformIds, Components components) {
	MutableNetwork<PipelineNode, PipelineEdge> network =
	NetworkBuilder.directed().allowsParallelEdges(true).allowsSelfLoops(false).build();
	Set<PCollectionNode> unproducedCollections = new HashSet<>();
	for (String transformId : transformIds) {
	PTransform transform = components.getTransformsOrThrow(transformId);
	PTransformNode transformNode =
	PipelineNode.pTransform(transformId, this.components.getTransformsOrThrow(transformId));
	network.addNode(transformNode);
	for (String produced : transform.getOutputsMap().values()) {
	PCollectionNode producedNode =
	PipelineNode.pCollection(produced, components.getPcollectionsOrThrow(produced));
	network.addNode(producedNode);
	network.addEdge(transformNode, producedNode, new PerElementEdge());
	checkArgument(
	network.inDegree(producedNode) == 1,
	"A %s should have exactly one producing %s, but found %s:\nPCollection:\n%s\nProducers:\n%s",
	PCollectionNode.class.getSimpleName(),
	PTransformNode.class.getSimpleName(),
	network.predecessors(producedNode).size(),
	producedNode,
	network.predecessors(producedNode));
	unproducedCollections.remove(producedNode);
	}
	for (Map.Entry<String, String> consumed : transform.getInputsMap().entrySet()) {
	// This loop may add an edge between the consumed PCollection and the current PTransform.
	// The local name of the transform must be used to determine the type of edge.
	String pcollectionId = consumed.getValue();
	PCollectionNode consumedNode =
	PipelineNode.pCollection(
	pcollectionId, this.components.getPcollectionsOrThrow(pcollectionId));
	if (network.addNode(consumedNode)) {
	// This node has been added to the network for the first time, so it has no producer.
	unproducedCollections.add(consumedNode);
	}
	if (getLocalSideInputNames(transform).contains(consumed.getKey())) {
	network.addEdge(consumedNode, transformNode, new SingletonEdge());
	} else {
	network.addEdge(consumedNode, transformNode, new PerElementEdge());
	}
	}
	}
	checkArgument(
	unproducedCollections.isEmpty(),
	"%ss %s were consumed but never produced",
	PCollectionNode.class.getSimpleName(),
	unproducedCollections);
	return network;
	}

	public Collection<PTransformNode> getTransforms() {
	return pipelineNetwork.nodes().stream()
	.filter(PTransformNode.class::isInstance)
	.map(PTransformNode.class::cast)
	.collect(Collectors.toList());
	}

	public Iterable<PTransformNode> getTopologicallyOrderedTransforms() {
	return StreamSupport.stream(
	Networks.topologicalOrder(pipelineNetwork, Comparator.comparing(PipelineNode::getId))
	.spliterator(),
	false)
	.filter(PTransformNode.class::isInstance)
	.map(PTransformNode.class::cast)
	.collect(Collectors.toList());
	}

	/**
	* Get the transforms that are roots of this {@link QueryablePipeline}. These are all nodes which
	* have no input {@link PCollection}.
	*/
	public Set<PTransformNode> getRootTransforms() {
	return pipelineNetwork.nodes().stream()
	.filter(pipelineNode -> pipelineNetwork.inEdges(pipelineNode).isEmpty())
	.map(pipelineNode -> (PTransformNode) pipelineNode)
	.collect(Collectors.toSet());
	}

	public PTransformNode getProducer(PCollectionNode pcollection) {
	return (PTransformNode) Iterables.getOnlyElement(pipelineNetwork.predecessors(pcollection));
	}

	/**
	* Get all of the {@link PTransformNode PTransforms} which consume the provided {@link
	* PCollectionNode} on a per-element basis.
	*
	* <p>If a {@link PTransformNode} consumes a {@link PCollectionNode} on a per-element basis one or
	* more times, it will appear a single time in the result.
	*
	* <p>In theory, a transform may consume a single {@link PCollectionNode} in both a per-element
	* and singleton manner. If this is the case, the transform node is included in the result, as it
	* does consume the {@link PCollectionNode} on a per-element basis.
	*/
	public Set<PTransformNode> getPerElementConsumers(PCollectionNode pCollection) {
	return pipelineNetwork.successors(pCollection).stream()
	.filter(
	consumer ->
	pipelineNetwork.edgesConnecting(pCollection, consumer).stream()
	.anyMatch(PipelineEdge::isPerElement))
	.map(pipelineNode -> (PTransformNode) pipelineNode)
	.collect(Collectors.toSet());
	}

	/**
	* Same as {@link #getPerElementConsumers(PCollectionNode)}, but returns transforms that consume
	* the collection as a singleton.
	*/
	public Set<PTransformNode> getSingletonConsumers(PCollectionNode pCollection) {
	return pipelineNetwork.successors(pCollection).stream()
	.filter(
	consumer ->
	pipelineNetwork.edgesConnecting(pCollection, consumer).stream()
	.anyMatch(edge -> !edge.isPerElement()))
	.map(pipelineNode -> (PTransformNode) pipelineNode)
	.collect(Collectors.toSet());
	}

	/**
	* Gets each {@link PCollectionNode} that the provided {@link PTransformNode} consumes on a
	* per-element basis.
	*/
	public Set<PCollectionNode> getPerElementInputPCollections(PTransformNode ptransform) {
	return pipelineNetwork.inEdges(ptransform).stream()
	.filter(PipelineEdge::isPerElement)
	.map(edge -> (PCollectionNode) pipelineNetwork.incidentNodes(edge).source())
	.collect(Collectors.toSet());
	}

	public Set<PCollectionNode> getOutputPCollections(PTransformNode ptransform) {
	return pipelineNetwork.successors(ptransform).stream()
	.map(pipelineNode -> (PCollectionNode) pipelineNode)
	.collect(Collectors.toSet());
	}

	public Components getComponents() {
	return components;
	}

	/**
	* Returns the {@link SideInputReference SideInputReferences} that the provided transform consumes
	* as side inputs.
	*/
	public Collection<SideInputReference> getSideInputs(PTransformNode transform) {
	return getLocalSideInputNames(transform.getTransform()).stream()
	.map(
	localName -> {
	String transformId = transform.getId();
	PTransform transformProto = components.getTransformsOrThrow(transformId);
	String collectionId = transform.getTransform().getInputsOrThrow(localName);
	PCollection collection = components.getPcollectionsOrThrow(collectionId);
	return SideInputReference.of(
	PipelineNode.pTransform(transformId, transformProto),
	localName,
	PipelineNode.pCollection(collectionId, collection));
	})
	.collect(Collectors.toSet());
	}

	public Collection<UserStateReference> getUserStates(PTransformNode transform) {
	return getLocalUserStateNames(transform.getTransform()).stream()
	.map(
	localName -> {
	String transformId = transform.getId();
	PTransform transformProto = components.getTransformsOrThrow(transformId);
	// Get the main input PCollection id.
	String collectionId =
	transform
	.getTransform()
	.getInputsOrThrow(
	Iterables.getOnlyElement(
	Sets.difference(
	transform.getTransform().getInputsMap().keySet(),
	ImmutableSet.builder()
	.addAll(getLocalSideInputNames(transformProto))
	.addAll(getLocalTimerNames(transformProto))
	.build())));
	PCollection collection = components.getPcollectionsOrThrow(collectionId);
	return UserStateReference.of(
	PipelineNode.pTransform(transformId, transformProto),
	localName,
	PipelineNode.pCollection(collectionId, collection));
	})
	.collect(Collectors.toSet());
	}

	public Collection<TimerReference> getTimers(PTransformNode transform) {
	return getLocalTimerNames(transform.getTransform()).stream()
	.map(
	localName -> {
	String transformId = transform.getId();
	PTransform transformProto = components.getTransformsOrThrow(transformId);
	return TimerReference.of(
	PipelineNode.pTransform(transformId, transformProto), localName);
	})
	.collect(Collectors.toSet());
	}

	private Set<String> getLocalSideInputNames(PTransform transform) {
	if (PAR_DO_TRANSFORM_URN.equals(transform.getSpec().getUrn())) {
	try {
	return ParDoPayload.parseFrom(transform.getSpec().getPayload()).getSideInputsMap().keySet();
	} catch (InvalidProtocolBufferException e) {
	throw new RuntimeException(e);
	}
	} else {
	return Collections.emptySet();
	}
	}

	private Set<String> getLocalUserStateNames(PTransform transform) {
	if (PAR_DO_TRANSFORM_URN.equals(transform.getSpec().getUrn())) {
	try {
	return ParDoPayload.parseFrom(transform.getSpec().getPayload()).getStateSpecsMap().keySet();
	} catch (InvalidProtocolBufferException e) {
	throw new RuntimeException(e);
	}
	} else {
	return Collections.emptySet();
	}
	}

	private Set<String> getLocalTimerNames(PTransform transform) {
	if (PAR_DO_TRANSFORM_URN.equals(transform.getSpec().getUrn())) {
	try {
	return ParDoPayload.parseFrom(transform.getSpec().getPayload()).getTimerSpecsMap().keySet();
	} catch (InvalidProtocolBufferException e) {
	throw new RuntimeException(e);
	}
	} else {
	return Collections.emptySet();
	}
	}

	public Optional<Environment> getEnvironment(PTransformNode parDo) {
	return Environments.getEnvironment(parDo.getId(), components);
	}

	private interface PipelineEdge {
	boolean isPerElement();
	}

	private static class PerElementEdge implements PipelineEdge {
	@Override
	public boolean isPerElement() {
	return true;
	}
	}

	private static class SingletonEdge implements PipelineEdge {
	@Override
	public boolean isPerElement() {
	return false;
	}
	}
	}