runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SdkComponents.java

/*
 * 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;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;

import com.google.common.base.Equivalence;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.BiMap;
import com.google.common.collect.HashBiMap;
import com.google.common.collect.ListMultimap;
import java.io.IOException;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.Pipeline.PipelineVisitor;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.common.runner.v1.RunnerApi;
import org.apache.beam.sdk.common.runner.v1.RunnerApi.Components;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.runners.TransformHierarchy.Node;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.util.NameUtils;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.WindowingStrategy;

/** SDK objects that will be represented at some later point within a {@link Components} object. */
class SdkComponents {
  private final RunnerApi.Components.Builder componentsBuilder;

  private final BiMap<AppliedPTransform<?, ?, ?>, String> transformIds;
  private final BiMap<PCollection<?>, String> pCollectionIds;
  private final BiMap<WindowingStrategy<?, ?>, String> windowingStrategyIds;

  /** A map of Coder to IDs. Coders are stored here with identity equivalence. */
  private final BiMap<Equivalence.Wrapper<? extends Coder<?>>, String> coderIds;
  // TODO: Specify environments

  /** Create a new {@link SdkComponents} with no components. */
  static SdkComponents create() {
    return new SdkComponents();
  }

  public static RunnerApi.Pipeline translatePipeline(Pipeline p) {
    final SdkComponents components = create();
    final Collection<String> rootIds = new HashSet<>();
    p.traverseTopologically(
        new PipelineVisitor.Defaults() {
          private final ListMultimap<Node, AppliedPTransform<?, ?, ?>> children =
              ArrayListMultimap.create();

          @Override
          public void leaveCompositeTransform(Node node) {
            if (node.isRootNode()) {
              for (AppliedPTransform<?, ?, ?> pipelineRoot : children.get(node)) {
                rootIds.add(components.getExistingPTransformId(pipelineRoot));
              }
            } else {
              children.put(node.getEnclosingNode(), node.toAppliedPTransform());
              try {
                components.registerPTransform(node.toAppliedPTransform(), children.get(node));
              } catch (IOException e) {
                throw new RuntimeException(e);
              }
            }
          }

          @Override
          public void visitPrimitiveTransform(Node node) {
            children.put(node.getEnclosingNode(), node.toAppliedPTransform());
            try {
              components.registerPTransform(
                  node.toAppliedPTransform(), Collections.<AppliedPTransform<?, ?, ?>>emptyList());
            } catch (IOException e) {
              throw new IllegalStateException(e);
            }
          }
        });
    // TODO: Display Data
    return RunnerApi.Pipeline.newBuilder()
        .setComponents(components.toComponents())
        .addAllRootTransformIds(rootIds)
        .build();
  }

  private SdkComponents() {
    this.componentsBuilder = RunnerApi.Components.newBuilder();
    this.transformIds = HashBiMap.create();
    this.pCollectionIds = HashBiMap.create();
    this.windowingStrategyIds = HashBiMap.create();
    this.coderIds = HashBiMap.create();
  }

  /**
   * Registers the provided {@link AppliedPTransform} into this {@link SdkComponents}, returning a
   * unique ID for the {@link AppliedPTransform}. Multiple registrations of the same
   * {@link AppliedPTransform} will return the same unique ID.
   *
   * <p>All of the children must already be registered within this {@link SdkComponents}.
   */
  String registerPTransform(
      AppliedPTransform<?, ?, ?> appliedPTransform, List<AppliedPTransform<?, ?, ?>> children)
      throws IOException {
    String name = getApplicationName(appliedPTransform);
    // If this transform is present in the components, nothing to do. return the existing name.
    // Otherwise the transform must be translated and added to the components.
    if (componentsBuilder.getTransformsOrDefault(name, null) != null) {
      return name;
    }
    checkNotNull(children, "child nodes may not be null");
    componentsBuilder.putTransforms(name, PTransforms.toProto(appliedPTransform, children, this));
    return name;
  }

  /**
   * Gets the ID for the provided {@link AppliedPTransform}. The provided {@link AppliedPTransform}
   * will not be added to the components produced by this {@link SdkComponents} until it is
   * translated via {@link #registerPTransform(AppliedPTransform, List)}.
   */
  private String getApplicationName(AppliedPTransform<?, ?, ?> appliedPTransform) {
    String existing = transformIds.get(appliedPTransform);
    if (existing != null) {
      return existing;
    }

    String name = appliedPTransform.getFullName();
    if (name.isEmpty()) {
      name = "unnamed-ptransform";
    }
    name = uniqify(name, transformIds.values());
    transformIds.put(appliedPTransform, name);
    return name;
  }

  String getExistingPTransformId(AppliedPTransform<?, ?, ?> appliedPTransform) {
    checkArgument(
        transformIds.containsKey(appliedPTransform),
        "%s %s has not been previously registered",
        AppliedPTransform.class.getSimpleName(),
        appliedPTransform);
    return transformIds.get(appliedPTransform);
  }

  /**
   * Registers the provided {@link PCollection} into this {@link SdkComponents}, returning a unique
   * ID for the {@link PCollection}. Multiple registrations of the same {@link PCollection} will
   * return the same unique ID.
   */
  String registerPCollection(PCollection<?> pCollection) throws IOException {
    String existing = pCollectionIds.get(pCollection);
    if (existing != null) {
      return existing;
    }
    String uniqueName = uniqify(pCollection.getName(), pCollectionIds.values());
    pCollectionIds.put(pCollection, uniqueName);
    componentsBuilder.putPcollections(uniqueName, PCollections.toProto(pCollection, this));
    return uniqueName;
  }

  /**
   * Registers the provided {@link WindowingStrategy} into this {@link SdkComponents}, returning a
   * unique ID for the {@link WindowingStrategy}. Multiple registrations of the same {@link
   * WindowingStrategy} will return the same unique ID.
   */
  String registerWindowingStrategy(WindowingStrategy<?, ?> windowingStrategy) throws IOException {
    String existing = windowingStrategyIds.get(windowingStrategy);
    if (existing != null) {
      return existing;
    }
    String baseName =
        String.format(
            "%s(%s)",
            NameUtils.approximateSimpleName(windowingStrategy),
            NameUtils.approximateSimpleName(windowingStrategy.getWindowFn()));
    String name = uniqify(baseName, windowingStrategyIds.values());
    windowingStrategyIds.put(windowingStrategy, name);
    RunnerApi.WindowingStrategy windowingStrategyProto =
        WindowingStrategies.toProto(windowingStrategy, this);
    componentsBuilder.putWindowingStrategies(name, windowingStrategyProto);
    return name;
  }

  /**
   * Registers the provided {@link Coder} into this {@link SdkComponents}, returning a unique ID for
   * the {@link Coder}. Multiple registrations of the same {@link Coder} will return the same
   * unique ID.
   *
   * <p>Coders are stored by identity to ensure that coders with implementations of {@link
   * #equals(Object)} and {@link #hashCode()} but incompatible binary formats are not considered the
   * same coder.
   */
  String registerCoder(Coder<?> coder) throws IOException {
    String existing = coderIds.get(Equivalence.identity().wrap(coder));
    if (existing != null) {
      return existing;
    }
    String baseName = NameUtils.approximateSimpleName(coder);
    String name = uniqify(baseName, coderIds.values());
    coderIds.put(Equivalence.identity().wrap(coder), name);
    RunnerApi.Coder coderProto = Coders.toProto(coder, this);
    componentsBuilder.putCoders(name, coderProto);
    return name;
  }

  private String uniqify(String baseName, Set<String> existing) {
    String name = baseName;
    int increment = 1;
    while (existing.contains(name)) {
      name = baseName + Integer.toString(increment);
      increment++;
    }
    return name;
  }

  /**
   * Convert this {@link SdkComponents} into a {@link RunnerApi.Components}, including all of the
   * contained {@link Coder coders}, {@link WindowingStrategy windowing strategies}, {@link
   * PCollection PCollections}, and {@link PTransform PTransforms}.
   */
  @Experimental
  RunnerApi.Components toComponents() {
    return componentsBuilder.build();
  }
}