/*
 * 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.direct.portable;

import static org.apache.beam.runners.core.construction.SyntheticComponents.uniqueId;
import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkArgument;
import static org.apache.beam.vendor.guava.v20_0.com.google.common.base.Preconditions.checkState;
import static org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables.getOnlyElement;

import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;
import org.apache.beam.model.fnexecution.v1.ProvisionApi.ProvisionInfo;
import org.apache.beam.model.fnexecution.v1.ProvisionApi.Resources;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.Coder;
import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
import org.apache.beam.model.pipeline.v1.RunnerApi.ComponentsOrBuilder;
import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec;
import org.apache.beam.model.pipeline.v1.RunnerApi.MessageWithComponents;
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.PTransform.Builder;
import org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline;
import org.apache.beam.model.pipeline.v1.RunnerApi.SdkFunctionSpec;
import org.apache.beam.runners.core.construction.ModelCoders;
import org.apache.beam.runners.core.construction.ModelCoders.KvCoderComponents;
import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.PipelineOptionsTranslation;
import org.apache.beam.runners.core.construction.graph.ExecutableStage;
import org.apache.beam.runners.core.construction.graph.GreedyPipelineFuser;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode;
import org.apache.beam.runners.core.construction.graph.PipelineValidator;
import org.apache.beam.runners.core.construction.graph.ProtoOverrides;
import org.apache.beam.runners.core.construction.graph.ProtoOverrides.TransformReplacement;
import org.apache.beam.runners.core.construction.graph.QueryablePipeline;
import org.apache.beam.runners.direct.ExecutableGraph;
import org.apache.beam.runners.fnexecution.GrpcContextHeaderAccessorProvider;
import org.apache.beam.runners.fnexecution.GrpcFnServer;
import org.apache.beam.runners.fnexecution.InProcessServerFactory;
import org.apache.beam.runners.fnexecution.ServerFactory;
import org.apache.beam.runners.fnexecution.artifact.ArtifactRetrievalService;
import org.apache.beam.runners.fnexecution.artifact.BeamFileSystemArtifactRetrievalService;
import org.apache.beam.runners.fnexecution.control.ControlClientPool;
import org.apache.beam.runners.fnexecution.control.FnApiControlClientPoolService;
import org.apache.beam.runners.fnexecution.control.JobBundleFactory;
import org.apache.beam.runners.fnexecution.control.MapControlClientPool;
import org.apache.beam.runners.fnexecution.control.SingleEnvironmentInstanceJobBundleFactory;
import org.apache.beam.runners.fnexecution.data.GrpcDataService;
import org.apache.beam.runners.fnexecution.environment.DockerEnvironmentFactory;
import org.apache.beam.runners.fnexecution.environment.EmbeddedEnvironmentFactory;
import org.apache.beam.runners.fnexecution.environment.EnvironmentFactory;
import org.apache.beam.runners.fnexecution.logging.GrpcLoggingService;
import org.apache.beam.runners.fnexecution.logging.Slf4jLogWriter;
import org.apache.beam.runners.fnexecution.provisioning.StaticGrpcProvisionService;
import org.apache.beam.runners.fnexecution.state.GrpcStateService;
import org.apache.beam.runners.fnexecution.wire.LengthPrefixUnknownCoders;
import org.apache.beam.sdk.fn.IdGenerator;
import org.apache.beam.sdk.fn.IdGenerators;
import org.apache.beam.sdk.fn.stream.OutboundObserverFactory;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.vendor.grpc.v1p13p1.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v20_0.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Maps;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Sets;
import org.joda.time.Duration;
import org.joda.time.Instant;

/**
 * The "ReferenceRunner" engine implementation. The ReferenceRunner uses the portability framework
 * to execute a Pipeline on a single machine.
 */
public class ReferenceRunner {
  private final RunnerApi.Pipeline pipeline;
  private final Struct options;
  private final String artifactRetrievalToken;

  private final EnvironmentType environmentType;

  private final IdGenerator idGenerator = IdGenerators.incrementingLongs();

  /** @param environmentType The environment to use for the SDK Harness. */
  private ReferenceRunner(
      Pipeline p, Struct options, String artifactRetrievalToken, EnvironmentType environmentType) {
    this.pipeline = executable(p);
    this.options = options;
    this.environmentType = environmentType;
    this.artifactRetrievalToken = artifactRetrievalToken;
  }

  /**
   * Creates a "ReferenceRunner" engine for a single pipeline with a Dockerized SDK harness.
   *
   * @param p Pipeline being executed for this job.
   * @param options PipelineOptions for this job.
   * @param artifactRetrievalToken Token to retrieve artifacts that have been staged.
   */
  public static ReferenceRunner forPipeline(
      RunnerApi.Pipeline p, Struct options, String artifactRetrievalToken) {
    return new ReferenceRunner(p, options, artifactRetrievalToken, EnvironmentType.DOCKER);
  }

  static ReferenceRunner forInProcessPipeline(RunnerApi.Pipeline p, Struct options) {
    return new ReferenceRunner(p, options, "", EnvironmentType.IN_PROCESS);
  }

  private RunnerApi.Pipeline executable(RunnerApi.Pipeline original) {
    RunnerApi.Pipeline p = original;
    PipelineValidator.validate(p);
    p =
        ProtoOverrides.updateTransform(
            PTransformTranslation.SPLITTABLE_PROCESS_KEYED_URN,
            p,
            new SplittableProcessKeyedReplacer());
    p =
        ProtoOverrides.updateTransform(
            PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, p, new PortableGroupByKeyReplacer());
    p = GreedyPipelineFuser.fuse(p).toPipeline();

    p = foldFeedSDFIntoExecutableStage(p);
    PipelineValidator.validate(p);

    return p;
  }

  private static Set<PCollectionNode> getKeyedPCollections(
      ExecutableGraph<PTransformNode, PCollectionNode> graph) {
    // This mimics KeyedPValueTrackingVisitor behavior in regular direct runner,
    // but without propagating keyed-ness through key-preserving DoFn's.
    // That is not yet necessary, but will be necessary once we implement state and timers.
    // See https://issues.apache.org/jira/browse/BEAM-4557.
    Set<PCollectionNode> res = Sets.newHashSet();
    Set<String> keyedProducers =
        Sets.newHashSet(DirectGroupByKey.DIRECT_GBKO_URN, DirectGroupByKey.DIRECT_GABW_URN);
    for (PTransformNode transform : graph.getExecutables()) {
      if (keyedProducers.contains(transform.getTransform().getSpec().getUrn())) {
        res.addAll(graph.getProduced(transform));
      }
    }
    return res;
  }

  /**
   * First starts all the services needed, then configures and starts the {@link
   * ExecutorServiceParallelExecutor}.
   */
  public void execute() throws Exception {
    ExecutableGraph<PTransformNode, PCollectionNode> graph = PortableGraph.forPipeline(pipeline);
    BundleFactory bundleFactory = ImmutableListBundleFactory.create();
    EvaluationContext ctxt =
        EvaluationContext.create(Instant::new, bundleFactory, graph, getKeyedPCollections(graph));
    RootProviderRegistry rootRegistry = RootProviderRegistry.javaPortableRegistry(bundleFactory);
    int targetParallelism = Math.max(Runtime.getRuntime().availableProcessors(), 3);
    ServerFactory serverFactory = createServerFactory();
    ControlClientPool controlClientPool = MapControlClientPool.create();
    ExecutorService dataExecutor = Executors.newCachedThreadPool();
    ProvisionInfo provisionInfo =
        ProvisionInfo.newBuilder()
            .setJobId("id")
            .setJobName("reference")
            .setPipelineOptions(options)
            .setWorkerId("foo")
            .setResourceLimits(Resources.getDefaultInstance())
            .setRetrievalToken(artifactRetrievalToken)
            .build();
    try (GrpcFnServer<GrpcLoggingService> logging =
            GrpcFnServer.allocatePortAndCreateFor(
                GrpcLoggingService.forWriter(Slf4jLogWriter.getDefault()), serverFactory);
        GrpcFnServer<ArtifactRetrievalService> artifact =
            GrpcFnServer.allocatePortAndCreateFor(
                BeamFileSystemArtifactRetrievalService.create(), serverFactory);
        GrpcFnServer<StaticGrpcProvisionService> provisioning =
            GrpcFnServer.allocatePortAndCreateFor(
                StaticGrpcProvisionService.create(provisionInfo), serverFactory);
        GrpcFnServer<FnApiControlClientPoolService> control =
            GrpcFnServer.allocatePortAndCreateFor(
                FnApiControlClientPoolService.offeringClientsToPool(
                    controlClientPool.getSink(),
                    GrpcContextHeaderAccessorProvider.getHeaderAccessor()),
                serverFactory);
        GrpcFnServer<GrpcDataService> data =
            GrpcFnServer.allocatePortAndCreateFor(
                GrpcDataService.create(dataExecutor, OutboundObserverFactory.serverDirect()),
                serverFactory);
        GrpcFnServer<GrpcStateService> state =
            GrpcFnServer.allocatePortAndCreateFor(GrpcStateService.create(), serverFactory)) {

      EnvironmentFactory environmentFactory =
          createEnvironmentFactory(control, logging, artifact, provisioning, controlClientPool);
      JobBundleFactory jobBundleFactory =
          SingleEnvironmentInstanceJobBundleFactory.create(
              environmentFactory, data, state, idGenerator);

      TransformEvaluatorRegistry transformRegistry =
          TransformEvaluatorRegistry.portableRegistry(
              graph,
              pipeline.getComponents(),
              bundleFactory,
              jobBundleFactory,
              EvaluationContextStepStateAndTimersProvider.forContext(ctxt));
      ExecutorServiceParallelExecutor executor =
          ExecutorServiceParallelExecutor.create(
              targetParallelism, rootRegistry, transformRegistry, graph, ctxt);
      executor.start();
      executor.waitUntilFinish(Duration.ZERO);
    } finally {
      dataExecutor.shutdown();
    }
  }

  private ServerFactory createServerFactory() {
    switch (environmentType) {
      case DOCKER:
        return ServerFactory.createDefault();
      case IN_PROCESS:
        return InProcessServerFactory.create();
      default:
        throw new IllegalArgumentException(
            String.format("Unknown %s %s", EnvironmentType.class.getSimpleName(), environmentType));
    }
  }

  private EnvironmentFactory createEnvironmentFactory(
      GrpcFnServer<FnApiControlClientPoolService> control,
      GrpcFnServer<GrpcLoggingService> logging,
      GrpcFnServer<ArtifactRetrievalService> artifact,
      GrpcFnServer<StaticGrpcProvisionService> provisioning,
      ControlClientPool controlClient) {
    switch (environmentType) {
      case DOCKER:
        return new DockerEnvironmentFactory.Provider(PipelineOptionsTranslation.fromProto(options))
            .createEnvironmentFactory(
                control, logging, artifact, provisioning, controlClient, idGenerator);
      case IN_PROCESS:
        return EmbeddedEnvironmentFactory.create(
            PipelineOptionsFactory.create(), logging, control, controlClient.getSource());
      default:
        throw new IllegalArgumentException(
            String.format("Unknown %s %s", EnvironmentType.class.getSimpleName(), environmentType));
    }
  }

  @VisibleForTesting
  static class PortableGroupByKeyReplacer implements TransformReplacement {
    @Override
    public MessageWithComponents getReplacement(String gbkId, ComponentsOrBuilder components) {
      PTransform gbk = components.getTransformsOrThrow(gbkId);
      checkArgument(
          PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN.equals(gbk.getSpec().getUrn()),
          "URN must be %s, got %s",
          PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN,
          gbk.getSpec().getUrn());

      PTransform.Builder newTransform = gbk.toBuilder();
      Components.Builder newComponents = Components.newBuilder();
      String inputId = getOnlyElement(gbk.getInputsMap().values());

      // Add the GBKO transform
      String kwiCollectionId =
          uniqueId(String.format("%s.%s", inputId, "kwi"), components::containsPcollections);
      {
        PCollection input = components.getPcollectionsOrThrow(inputId);
        Coder inputCoder = components.getCodersOrThrow(input.getCoderId());
        KvCoderComponents kvComponents = ModelCoders.getKvCoderComponents(inputCoder);
        String windowCoderId =
            components
                .getWindowingStrategiesOrThrow(input.getWindowingStrategyId())
                .getWindowCoderId();
        // This coder isn't actually required for the pipeline to function properly - the KWIs can
        // be passed around as pure java objects with no coding of the values, but it approximates
        // a full pipeline.
        Coder kwiCoder =
            Coder.newBuilder()
                .setSpec(
                    SdkFunctionSpec.newBuilder()
                        .setSpec(FunctionSpec.newBuilder().setUrn("beam:direct:keyedworkitem:v1")))
                .addAllComponentCoderIds(
                    ImmutableList.of(
                        kvComponents.keyCoderId(), kvComponents.valueCoderId(), windowCoderId))
                .build();
        String kwiCoderId =
            uniqueId(
                String.format("kwi(%s:%s)", kvComponents.keyCoderId(), kvComponents.valueCoderId()),
                components::containsCoders);
        // The kwi PCollection has the same WindowingStrategy as the input, as no merging will
        // have been performed, so elements remain in their original windows
        PCollection kwi =
            input.toBuilder().setUniqueName(kwiCollectionId).setCoderId(kwiCoderId).build();
        String gbkoId = uniqueId(String.format("%s/GBKO", gbkId), components::containsTransforms);
        PTransform gbko =
            PTransform.newBuilder()
                .setUniqueName(String.format("%s/GBKO", gbk.getUniqueName()))
                .putAllInputs(gbk.getInputsMap())
                .setSpec(FunctionSpec.newBuilder().setUrn(DirectGroupByKey.DIRECT_GBKO_URN))
                .putOutputs("output", kwiCollectionId)
                .build();

        newTransform.addSubtransforms(gbkoId);
        newComponents
            .putCoders(kwiCoderId, kwiCoder)
            .putPcollections(kwiCollectionId, kwi)
            .putTransforms(gbkoId, gbko);
      }

      // Add the GABW transform
      {
        String gabwId = uniqueId(String.format("%s/GABW", gbkId), components::containsTransforms);
        PTransform gabw =
            PTransform.newBuilder()
                .setUniqueName(String.format("%s/GABW", gbk.getUniqueName()))
                .putInputs("input", kwiCollectionId)
                .setSpec(FunctionSpec.newBuilder().setUrn(DirectGroupByKey.DIRECT_GABW_URN))
                .putAllOutputs(gbk.getOutputsMap())
                .build();
        newTransform.addSubtransforms(gabwId);
        newComponents.putTransforms(gabwId, gabw);
      }

      return MessageWithComponents.newBuilder()
          .setPtransform(newTransform)
          .setComponents(newComponents)
          .build();
    }
  }

  /**
   * Replaces the {@link PTransformTranslation#SPLITTABLE_PROCESS_KEYED_URN} with a {@link
   * DirectGroupByKey#DIRECT_GBKO_URN} (construct keyed work items) followed by a {@link
   * SplittableRemoteStageEvaluatorFactory#FEED_SDF_URN} (convert the keyed work items to
   * element/restriction pairs that later go into {@link
   * PTransformTranslation#SPLITTABLE_PROCESS_ELEMENTS_URN}).
   */
  @VisibleForTesting
  static class SplittableProcessKeyedReplacer implements TransformReplacement {
    @Override
    public MessageWithComponents getReplacement(String spkId, ComponentsOrBuilder components) {
      PTransform spk = components.getTransformsOrThrow(spkId);
      checkArgument(
          PTransformTranslation.SPLITTABLE_PROCESS_KEYED_URN.equals(spk.getSpec().getUrn()),
          "URN must be %s, got %s",
          PTransformTranslation.SPLITTABLE_PROCESS_KEYED_URN,
          spk.getSpec().getUrn());

      Components.Builder newComponents = Components.newBuilder();
      newComponents.putAllCoders(components.getCodersMap());

      Builder newPTransform = spk.toBuilder();

      String inputId = getOnlyElement(spk.getInputsMap().values());
      PCollection input = components.getPcollectionsOrThrow(inputId);

      // This is a Coder<KV<String, KV<ElementT, RestrictionT>>>
      Coder inputCoder = components.getCodersOrThrow(input.getCoderId());
      KvCoderComponents kvComponents = ModelCoders.getKvCoderComponents(inputCoder);
      String windowCoderId =
          components
              .getWindowingStrategiesOrThrow(input.getWindowingStrategyId())
              .getWindowCoderId();

      // === Construct a raw GBK returning KeyedWorkItem's ===
      String kwiCollectionId =
          uniqueId(String.format("%s.kwi", spkId), components::containsPcollections);
      {
        // This coder isn't actually required for the pipeline to function properly - the KWIs can
        // be passed around as pure java objects with no coding of the values, but it approximates a
        // full pipeline.
        Coder kwiCoder =
            Coder.newBuilder()
                .setSpec(
                    SdkFunctionSpec.newBuilder()
                        .setSpec(FunctionSpec.newBuilder().setUrn("beam:direct:keyedworkitem:v1")))
                .addAllComponentCoderIds(
                    ImmutableList.of(
                        kvComponents.keyCoderId(), kvComponents.valueCoderId(), windowCoderId))
                .build();
        String kwiCoderId =
            uniqueId(
                String.format(
                    "keyed_work_item(%s:%s)",
                    kvComponents.keyCoderId(), kvComponents.valueCoderId()),
                components::containsCoders);

        PCollection kwiCollection =
            input.toBuilder().setUniqueName(kwiCollectionId).setCoderId(kwiCoderId).build();
        String rawGbkId =
            uniqueId(String.format("%s/RawGBK", spkId), components::containsTransforms);
        PTransform rawGbk =
            PTransform.newBuilder()
                .setUniqueName(String.format("%s/RawGBK", spk.getUniqueName()))
                .putAllInputs(spk.getInputsMap())
                .setSpec(FunctionSpec.newBuilder().setUrn(DirectGroupByKey.DIRECT_GBKO_URN))
                .putOutputs("output", kwiCollectionId)
                .build();

        newComponents
            .putCoders(kwiCoderId, kwiCoder)
            .putPcollections(kwiCollectionId, kwiCollection)
            .putTransforms(rawGbkId, rawGbk);
        newPTransform.addSubtransforms(rawGbkId);
      }

      // === Construct a "Feed SDF" operation that converts KWI to KV<ElementT, RestrictionT> ===
      String feedSDFCollectionId =
          uniqueId(String.format("%s.feed", spkId), components::containsPcollections);
      {
        String elementRestrictionCoderId = kvComponents.valueCoderId();
        String feedSDFCoderId =
            LengthPrefixUnknownCoders.addLengthPrefixedCoder(
                elementRestrictionCoderId, newComponents, false);

        PCollection feedSDFCollection =
            input.toBuilder().setUniqueName(feedSDFCollectionId).setCoderId(feedSDFCoderId).build();
        String feedSDFId =
            uniqueId(String.format("%s/FeedSDF", spkId), components::containsTransforms);
        PTransform feedSDF =
            PTransform.newBuilder()
                .setUniqueName(String.format("%s/FeedSDF", spk.getUniqueName()))
                .putInputs("input", kwiCollectionId)
                .setSpec(
                    FunctionSpec.newBuilder()
                        .setUrn(SplittableRemoteStageEvaluatorFactory.FEED_SDF_URN))
                .putOutputs("output", feedSDFCollectionId)
                .build();

        newComponents
            .putPcollections(feedSDFCollectionId, feedSDFCollection)
            .putTransforms(feedSDFId, feedSDF);
        newPTransform.addSubtransforms(feedSDFId);
      }

      // === Construct the SPLITTABLE_PROCESS_ELEMENTS operation
      {
        String runSDFId =
            uniqueId(String.format("%s/RunSDF", spkId), components::containsTransforms);
        PTransform runSDF =
            PTransform.newBuilder()
                .setUniqueName(String.format("%s/RunSDF", spk.getUniqueName()))
                .putInputs("input", feedSDFCollectionId)
                .setSpec(
                    FunctionSpec.newBuilder()
                        .setUrn(PTransformTranslation.SPLITTABLE_PROCESS_ELEMENTS_URN)
                        .setPayload(spk.getSpec().getPayload()))
                .putAllOutputs(spk.getOutputsMap())
                .build();
        newComponents.putTransforms(runSDFId, runSDF);
        newPTransform.addSubtransforms(runSDFId);
      }

      return MessageWithComponents.newBuilder()
          .setPtransform(newPTransform.build())
          .setComponents(newComponents)
          .build();
    }
  }

  /**
   * Finds FEED_SDF nodes followed by an ExecutableStage and replaces them by a single {@link
   * SplittableRemoteStageEvaluatorFactory#URN} stage that feeds the ExecutableStage knowing that
   * the first instruction in the stage is an SDF.
   */
  private static Pipeline foldFeedSDFIntoExecutableStage(Pipeline p) {
    Pipeline.Builder newPipeline = p.toBuilder();
    Components.Builder newPipelineComponents = newPipeline.getComponentsBuilder();

    QueryablePipeline q = QueryablePipeline.forPipeline(p);
    String feedSdfUrn = SplittableRemoteStageEvaluatorFactory.FEED_SDF_URN;
    List<PTransformNode> feedSDFNodes =
        q.getTransforms().stream()
            .filter(node -> node.getTransform().getSpec().getUrn().equals(feedSdfUrn))
            .collect(Collectors.toList());
    Map<String, PTransformNode> stageToFeeder = Maps.newHashMap();
    for (PTransformNode node : feedSDFNodes) {
      PCollectionNode output = Iterables.getOnlyElement(q.getOutputPCollections(node));
      PTransformNode consumer = Iterables.getOnlyElement(q.getPerElementConsumers(output));
      String consumerUrn = consumer.getTransform().getSpec().getUrn();
      checkState(
          consumerUrn.equals(ExecutableStage.URN),
          "Expected all FeedSDF nodes to be consumed by an ExecutableStage, "
              + "but %s is consumed by %s which is %s",
          node.getId(),
          consumer.getId(),
          consumerUrn);
      stageToFeeder.put(consumer.getId(), node);
    }

    // Copy over root transforms except for the excluded FEED_SDF transforms.
    Set<String> feedSDFIds =
        feedSDFNodes.stream().map(PTransformNode::getId).collect(Collectors.toSet());
    newPipeline.clearRootTransformIds();
    for (String rootId : p.getRootTransformIdsList()) {
      if (!feedSDFIds.contains(rootId)) {
        newPipeline.addRootTransformIds(rootId);
      }
    }
    // Copy over all transforms, except FEED_SDF transforms are skipped, and ExecutableStage's
    // feeding from them are replaced.
    for (PTransformNode node : q.getTransforms()) {
      if (feedSDFNodes.contains(node)) {
        // These transforms are skipped and handled separately.
        continue;
      }
      if (!stageToFeeder.containsKey(node.getId())) {
        // This transform is unchanged
        newPipelineComponents.putTransforms(node.getId(), node.getTransform());
        continue;
      }
      // "node" is an ExecutableStage transform feeding from an SDF.
      PTransformNode feedSDFNode = stageToFeeder.get(node.getId());
      PCollectionNode rawGBKOutput =
          Iterables.getOnlyElement(q.getPerElementInputPCollections(feedSDFNode));

      // Replace the ExecutableStage transform.
      newPipelineComponents.putTransforms(
          node.getId(),
          node.getTransform()
              .toBuilder()
              .mergeSpec(
                  // Change URN from ExecutableStage.URN to URN of the ULR's splittable executable
                  // stage evaluator.
                  FunctionSpec.newBuilder()
                      .setUrn(SplittableRemoteStageEvaluatorFactory.URN)
                      .build())
              .putInputs(
                  // The splittable executable stage now reads from the raw GBK, instead of
                  // from the now non-existent FEED_SDF.
                  Iterables.getOnlyElement(node.getTransform().getInputsMap().keySet()),
                  rawGBKOutput.getId())
              .build());
    }

    return newPipeline.build();
  }

  private enum EnvironmentType {
    DOCKER,
    IN_PROCESS
  }
}