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

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

import com.google.common.collect.BiMap;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
import java.util.stream.Collectors;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.SideInputId;
import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.RehydratedComponents;
import org.apache.beam.runners.core.construction.WindowingStrategyTranslation;
import org.apache.beam.runners.core.construction.graph.ExecutableStage;
import org.apache.beam.runners.core.construction.graph.PipelineNode;
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.QueryablePipeline;
import org.apache.beam.runners.flink.translation.functions.FlinkAssignWindows;
import org.apache.beam.runners.flink.translation.functions.FlinkExecutableStageContext;
import org.apache.beam.runners.flink.translation.functions.FlinkExecutableStageFunction;
import org.apache.beam.runners.flink.translation.functions.FlinkExecutableStagePruningFunction;
import org.apache.beam.runners.flink.translation.functions.FlinkPartialReduceFunction;
import org.apache.beam.runners.flink.translation.functions.FlinkReduceFunction;
import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
import org.apache.beam.runners.flink.translation.types.KvKeySelector;
import org.apache.beam.runners.flink.translation.utils.FlinkPipelineTranslatorUtils;
import org.apache.beam.runners.flink.translation.wrappers.ImpulseInputFormat;
import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
import org.apache.beam.runners.fnexecution.wire.WireCoders;
import org.apache.beam.sdk.coders.ByteArrayCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.ListCoder;
import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.join.RawUnionValue;
import org.apache.beam.sdk.transforms.join.UnionCoder;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.apache.beam.vendor.protobuf.v3.com.google.protobuf.InvalidProtocolBufferException;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.DataSet;
import org.apache.flink.api.java.ExecutionEnvironment;
import org.apache.flink.api.java.io.DiscardingOutputFormat;
import org.apache.flink.api.java.operators.DataSource;
import org.apache.flink.api.java.operators.FlatMapOperator;
import org.apache.flink.api.java.operators.GroupCombineOperator;
import org.apache.flink.api.java.operators.GroupReduceOperator;
import org.apache.flink.api.java.operators.Grouping;
import org.apache.flink.api.java.operators.MapPartitionOperator;

/**
 * A translator that translates bounded portable pipelines into executable Flink pipelines.
 *
 * <p>Example usage:
 *
 * <pre>
 *   FlinkBatchPortablePipelineTranslator translator =
 *       FlinkBatchPortablePipelineTranslator.createTranslator();
 *   BatchTranslationContext context =
 *       FlinkBatchPortablePipelineTranslator.createTranslationContext(jobInfo);
 *   translator.translate(context, pipeline);
 *   ExecutionEnvironment executionEnvironment = context.getExecutionEnvironment();
 *   // Do something with executionEnvironment...
 * </pre>
 *
 * <p>After translation the {@link ExecutionEnvironment} in the translation context will contain the
 * full not-yet-executed pipeline DAG corresponding to the input pipeline.
 */
public class FlinkBatchPortablePipelineTranslator
    implements FlinkPortablePipelineTranslator<
        FlinkBatchPortablePipelineTranslator.BatchTranslationContext> {

  /**
   * Creates a batch translation context. The resulting Flink execution dag will live in a new
   * {@link ExecutionEnvironment}.
   */
  public static BatchTranslationContext createTranslationContext(
      JobInfo jobInfo, FlinkPipelineOptions pipelineOptions, List<String> filesToStage) {
    ExecutionEnvironment executionEnvironment =
        FlinkExecutionEnvironments.createBatchExecutionEnvironment(pipelineOptions, filesToStage);
    return new BatchTranslationContext(jobInfo, pipelineOptions, executionEnvironment);
  }

  /** Creates a batch translator. */
  public static FlinkBatchPortablePipelineTranslator createTranslator() {
    ImmutableMap.Builder<String, PTransformTranslator> translatorMap = ImmutableMap.builder();
    translatorMap.put(
        PTransformTranslation.FLATTEN_TRANSFORM_URN,
        FlinkBatchPortablePipelineTranslator::translateFlatten);
    translatorMap.put(
        PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN,
        FlinkBatchPortablePipelineTranslator::translateGroupByKey);
    translatorMap.put(
        PTransformTranslation.IMPULSE_TRANSFORM_URN,
        FlinkBatchPortablePipelineTranslator::translateImpulse);
    translatorMap.put(
        PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN,
        FlinkBatchPortablePipelineTranslator::translateAssignWindows);
    translatorMap.put(
        ExecutableStage.URN, FlinkBatchPortablePipelineTranslator::translateExecutableStage);
    translatorMap.put(
        PTransformTranslation.RESHUFFLE_URN,
        FlinkBatchPortablePipelineTranslator::translateReshuffle);
    return new FlinkBatchPortablePipelineTranslator(translatorMap.build());
  }

  /**
   * Batch translation context. Stores metadata about known PCollections/DataSets and holds the
   * flink {@link ExecutionEnvironment} that the execution plan will be applied to.
   */
  public static class BatchTranslationContext
      implements FlinkPortablePipelineTranslator.TranslationContext {

    private final JobInfo jobInfo;
    private final FlinkPipelineOptions options;
    private final ExecutionEnvironment executionEnvironment;
    private final Map<String, DataSet<?>> dataSets;
    private final Set<String> danglingDataSets;

    private BatchTranslationContext(
        JobInfo jobInfo, FlinkPipelineOptions options, ExecutionEnvironment executionEnvironment) {
      this.jobInfo = jobInfo;
      this.options = options;
      this.executionEnvironment = executionEnvironment;
      dataSets = new HashMap<>();
      danglingDataSets = new HashSet<>();
    }

    @Override
    public JobInfo getJobInfo() {
      return jobInfo;
    }

    @Override
    public FlinkPipelineOptions getPipelineOptions() {
      return options;
    }

    public ExecutionEnvironment getExecutionEnvironment() {
      return executionEnvironment;
    }

    public <T> void addDataSet(String pCollectionId, DataSet<T> dataSet) {
      checkArgument(!dataSets.containsKey(pCollectionId));
      dataSets.put(pCollectionId, dataSet);
      danglingDataSets.add(pCollectionId);
    }

    public <T> DataSet<T> getDataSetOrThrow(String pCollectionId) {
      DataSet<T> dataSet = (DataSet<T>) dataSets.get(pCollectionId);
      if (dataSet == null) {
        throw new IllegalArgumentException(
            String.format("Unknown dataset for id %s.", pCollectionId));
      }

      // Assume that the DataSet is consumed if requested. We use this as a proxy for consumption
      // because Flink does not expose its internal execution plan.
      danglingDataSets.remove(pCollectionId);
      return dataSet;
    }

    public Collection<DataSet<?>> getDanglingDataSets() {
      return danglingDataSets.stream().map(id -> dataSets.get(id)).collect(Collectors.toList());
    }
  }

  /** Transform translation interface. */
  @FunctionalInterface
  private interface PTransformTranslator {
    /** Translate a PTransform into the given translation context. */
    void translate(
        PTransformNode transform, RunnerApi.Pipeline pipeline, BatchTranslationContext context);
  }

  private final Map<String, PTransformTranslator> urnToTransformTranslator;

  private FlinkBatchPortablePipelineTranslator(
      Map<String, PTransformTranslator> urnToTransformTranslator) {
    this.urnToTransformTranslator = urnToTransformTranslator;
  }

  @Override
  public void translate(BatchTranslationContext context, RunnerApi.Pipeline pipeline) {
    // Use a QueryablePipeline to traverse transforms topologically.
    QueryablePipeline p =
        QueryablePipeline.forTransforms(
            pipeline.getRootTransformIdsList(), pipeline.getComponents());
    for (PipelineNode.PTransformNode transform : p.getTopologicallyOrderedTransforms()) {
      urnToTransformTranslator
          .getOrDefault(
              transform.getTransform().getSpec().getUrn(),
              FlinkBatchPortablePipelineTranslator::urnNotFound)
          .translate(transform, pipeline, context);
    }

    // Ensure that side effects are performed for unconsumed DataSets.
    for (DataSet<?> dataSet : context.getDanglingDataSets()) {
      dataSet.output(new DiscardingOutputFormat<>());
    }
  }

  private static <K, V> void translateReshuffle(
      PTransformNode transform, RunnerApi.Pipeline pipeline, BatchTranslationContext context) {
    DataSet<WindowedValue<KV<K, V>>> inputDataSet =
        context.getDataSetOrThrow(
            Iterables.getOnlyElement(transform.getTransform().getInputsMap().values()));
    context.addDataSet(
        Iterables.getOnlyElement(transform.getTransform().getOutputsMap().values()),
        inputDataSet.rebalance());
  }

  private static <T> void translateAssignWindows(
      PTransformNode transform, RunnerApi.Pipeline pipeline, BatchTranslationContext context) {
    RunnerApi.Components components = pipeline.getComponents();
    RunnerApi.WindowIntoPayload payload;
    try {
      payload =
          RunnerApi.WindowIntoPayload.parseFrom(transform.getTransform().getSpec().getPayload());
    } catch (InvalidProtocolBufferException e) {
      throw new IllegalArgumentException(e);
    }
    //TODO: https://issues.apache.org/jira/browse/BEAM-4296
    // This only works for well known window fns, we should defer this execution to the SDK
    // if the WindowFn can't be parsed or just defer it all the time.
    WindowFn<T, ? extends BoundedWindow> windowFn =
        (WindowFn<T, ? extends BoundedWindow>)
            WindowingStrategyTranslation.windowFnFromProto(payload.getWindowFn());

    String inputCollectionId =
        Iterables.getOnlyElement(transform.getTransform().getInputsMap().values());
    String outputCollectionId =
        Iterables.getOnlyElement(transform.getTransform().getOutputsMap().values());
    PCollectionNode collectionNode =
        PipelineNode.pCollection(
            outputCollectionId, components.getPcollectionsOrThrow(outputCollectionId));
    Coder<WindowedValue<T>> outputCoder;
    try {
      outputCoder = WireCoders.instantiateRunnerWireCoder(collectionNode, components);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    TypeInformation<WindowedValue<T>> resultTypeInfo = new CoderTypeInformation<>(outputCoder);

    DataSet<WindowedValue<T>> inputDataSet = context.getDataSetOrThrow(inputCollectionId);

    FlinkAssignWindows<T, ? extends BoundedWindow> assignWindowsFunction =
        new FlinkAssignWindows<>(windowFn);

    DataSet<WindowedValue<T>> resultDataSet =
        inputDataSet
            .flatMap(assignWindowsFunction)
            .name(transform.getTransform().getUniqueName())
            .returns(resultTypeInfo);

    context.addDataSet(outputCollectionId, resultDataSet);
  }

  private static <InputT> void translateExecutableStage(
      PTransformNode transform, RunnerApi.Pipeline pipeline, BatchTranslationContext context) {
    // TODO: Fail on stateful DoFns for now.
    // TODO: Support stateful DoFns by inserting group-by-keys where necessary.
    // TODO: Fail on splittable DoFns.
    // TODO: Special-case single outputs to avoid multiplexing PCollections.

    RunnerApi.Components components = pipeline.getComponents();
    Map<String, String> outputs = transform.getTransform().getOutputsMap();
    // Mapping from PCollection id to coder tag id.
    BiMap<String, Integer> outputMap =
        FlinkPipelineTranslatorUtils.createOutputMap(outputs.values());
    // Collect all output Coders and create a UnionCoder for our tagged outputs.
    List<Coder<?>> unionCoders = Lists.newArrayList();
    // Enforce tuple tag sorting by union tag index.
    Map<String, Coder<WindowedValue<?>>> outputCoders = Maps.newHashMap();
    for (String collectionId : new TreeMap<>(outputMap.inverse()).values()) {
      PCollectionNode collectionNode =
          PipelineNode.pCollection(collectionId, components.getPcollectionsOrThrow(collectionId));
      Coder<WindowedValue<?>> coder;
      try {
        coder = (Coder) WireCoders.instantiateRunnerWireCoder(collectionNode, components);
      } catch (IOException e) {
        throw new RuntimeException(e);
      }
      outputCoders.put(collectionId, coder);
      unionCoders.add(coder);
    }
    UnionCoder unionCoder = UnionCoder.of(unionCoders);
    TypeInformation<RawUnionValue> typeInformation = new CoderTypeInformation<>(unionCoder);

    RunnerApi.ExecutableStagePayload stagePayload;
    try {
      stagePayload =
          RunnerApi.ExecutableStagePayload.parseFrom(
              transform.getTransform().getSpec().getPayload());
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    FlinkExecutableStageFunction<InputT> function =
        new FlinkExecutableStageFunction<>(
            stagePayload,
            context.getJobInfo(),
            outputMap,
            FlinkExecutableStageContext.factory(context.getPipelineOptions()));

    DataSet<WindowedValue<InputT>> inputDataSet =
        context.getDataSetOrThrow(stagePayload.getInput());

    MapPartitionOperator<WindowedValue<InputT>, RawUnionValue> taggedDataset =
        new MapPartitionOperator<>(
            inputDataSet, typeInformation, function, transform.getTransform().getUniqueName());

    for (SideInputId sideInputId : stagePayload.getSideInputsList()) {
      String collectionId =
          components
              .getTransformsOrThrow(sideInputId.getTransformId())
              .getInputsOrThrow(sideInputId.getLocalName());
      // Register under the global PCollection name. Only ExecutableStageFunction needs to know the
      // mapping from local name to global name and how to translate the broadcast data to a state
      // API view.
      taggedDataset.withBroadcastSet(context.getDataSetOrThrow(collectionId), collectionId);
    }

    for (String collectionId : outputs.values()) {
      pruneOutput(
          taggedDataset,
          context,
          outputMap.get(collectionId),
          outputCoders.get(collectionId),
          transform.getTransform().getUniqueName(),
          collectionId);
    }
    if (outputs.isEmpty()) {
      // NOTE: After pipeline translation, we traverse the set of unconsumed PCollections and add a
      // no-op sink to each to make sure they are materialized by Flink. However, some SDK-executed
      // stages have no runner-visible output after fusion. We handle this case by adding a sink
      // here.
      taggedDataset.output(new DiscardingOutputFormat());
    }
  }

  private static <T> void translateFlatten(
      PTransformNode transform, RunnerApi.Pipeline pipeline, BatchTranslationContext context) {
    Map<String, String> allInputs = transform.getTransform().getInputsMap();
    DataSet<WindowedValue<T>> result = null;

    if (allInputs.isEmpty()) {

      // Create an empty dummy source to satisfy downstream operations. We cannot create an empty
      // source in Flink, so we send the DataSet to a flatMap that never forwards its element.
      DataSource<String> dummySource = context.getExecutionEnvironment().fromElements("dummy");
      result =
          dummySource
              .<WindowedValue<T>>flatMap(
                  (s, collector) -> {
                    // never return anything
                  })
              .returns(
                  new CoderTypeInformation<>(
                      WindowedValue.getFullCoder(
                          (Coder<T>) VoidCoder.of(), GlobalWindow.Coder.INSTANCE)));
    } else {
      for (String pCollectionId : allInputs.values()) {
        DataSet<WindowedValue<T>> current = context.getDataSetOrThrow(pCollectionId);
        if (result == null) {
          result = current;
        } else {
          result = result.union(current);
        }
      }
    }

    // Insert a dummy filter. Flink produces duplicate elements after the union in some cases if we
    // don't do so.
    result = result.filter(tWindowedValue -> true).name("UnionFixFilter");
    context.addDataSet(
        Iterables.getOnlyElement(transform.getTransform().getOutputsMap().values()), result);
  }

  private static <K, V> void translateGroupByKey(
      PTransformNode transform, RunnerApi.Pipeline pipeline, BatchTranslationContext context) {
    RunnerApi.Components components = pipeline.getComponents();
    String inputPCollectionId =
        Iterables.getOnlyElement(transform.getTransform().getInputsMap().values());
    PCollectionNode inputCollection =
        PipelineNode.pCollection(
            inputPCollectionId, components.getPcollectionsOrThrow(inputPCollectionId));
    DataSet<WindowedValue<KV<K, V>>> inputDataSet = context.getDataSetOrThrow(inputPCollectionId);
    RunnerApi.WindowingStrategy windowingStrategyProto =
        pipeline
            .getComponents()
            .getWindowingStrategiesOrThrow(
                pipeline
                    .getComponents()
                    .getPcollectionsOrThrow(inputPCollectionId)
                    .getWindowingStrategyId());

    RehydratedComponents rehydratedComponents =
        RehydratedComponents.forComponents(pipeline.getComponents());

    WindowingStrategy<Object, BoundedWindow> windowingStrategy;
    try {
      windowingStrategy =
          (WindowingStrategy<Object, BoundedWindow>)
              WindowingStrategyTranslation.fromProto(windowingStrategyProto, rehydratedComponents);
    } catch (InvalidProtocolBufferException e) {
      throw new IllegalStateException(
          String.format(
              "Unable to hydrate GroupByKey windowing strategy %s.", windowingStrategyProto),
          e);
    }

    WindowedValueCoder<KV<K, V>> inputCoder;
    try {
      inputCoder =
          (WindowedValueCoder)
              WireCoders.instantiateRunnerWireCoder(inputCollection, pipeline.getComponents());
    } catch (IOException e) {
      throw new RuntimeException(e);
    }

    KvCoder<K, V> inputElementCoder = (KvCoder<K, V>) inputCoder.getValueCoder();

    Concatenate<V> combineFn = new Concatenate<>();
    Coder<List<V>> accumulatorCoder =
        combineFn.getAccumulatorCoder(
            CoderRegistry.createDefault(), inputElementCoder.getValueCoder());

    Coder<WindowedValue<KV<K, List<V>>>> outputCoder =
        WindowedValue.getFullCoder(
            KvCoder.of(inputElementCoder.getKeyCoder(), accumulatorCoder),
            windowingStrategy.getWindowFn().windowCoder());

    TypeInformation<WindowedValue<KV<K, List<V>>>> partialReduceTypeInfo =
        new CoderTypeInformation<>(outputCoder);

    Grouping<WindowedValue<KV<K, V>>> inputGrouping =
        inputDataSet.groupBy(new KvKeySelector<>(inputElementCoder.getKeyCoder()));

    FlinkPartialReduceFunction<K, V, List<V>, ?> partialReduceFunction =
        new FlinkPartialReduceFunction<>(
            combineFn, windowingStrategy, Collections.emptyMap(), context.getPipelineOptions());

    FlinkReduceFunction<K, List<V>, List<V>, ?> reduceFunction =
        new FlinkReduceFunction<>(
            combineFn, windowingStrategy, Collections.emptyMap(), context.getPipelineOptions());

    // Partially GroupReduce the values into the intermediate format AccumT (combine)
    GroupCombineOperator<WindowedValue<KV<K, V>>, WindowedValue<KV<K, List<V>>>> groupCombine =
        new GroupCombineOperator<>(
            inputGrouping,
            partialReduceTypeInfo,
            partialReduceFunction,
            "GroupCombine: " + transform.getTransform().getUniqueName());

    Grouping<WindowedValue<KV<K, List<V>>>> intermediateGrouping =
        groupCombine.groupBy(new KvKeySelector<>(inputElementCoder.getKeyCoder()));

    // Fully reduce the values and create output format VO
    GroupReduceOperator<WindowedValue<KV<K, List<V>>>, WindowedValue<KV<K, List<V>>>>
        outputDataSet =
            new GroupReduceOperator<>(
                intermediateGrouping,
                partialReduceTypeInfo,
                reduceFunction,
                transform.getTransform().getUniqueName());

    context.addDataSet(
        Iterables.getOnlyElement(transform.getTransform().getOutputsMap().values()), outputDataSet);
  }

  private static void translateImpulse(
      PTransformNode transform, RunnerApi.Pipeline pipeline, BatchTranslationContext context) {
    TypeInformation<WindowedValue<byte[]>> typeInformation =
        new CoderTypeInformation<>(
            WindowedValue.getFullCoder(ByteArrayCoder.of(), GlobalWindow.Coder.INSTANCE));
    DataSource<WindowedValue<byte[]>> dataSource =
        new DataSource<>(
            context.getExecutionEnvironment(),
            new ImpulseInputFormat(),
            typeInformation,
            transform.getTransform().getUniqueName());

    context.addDataSet(
        Iterables.getOnlyElement(transform.getTransform().getOutputsMap().values()), dataSource);
  }

  /**
   * Combiner that combines {@code T}s into a single {@code List<T>} containing all inputs.
   *
   * <p>For internal use to translate {@link GroupByKey}. For a large {@link PCollection} this is
   * expected to crash!
   *
   * <p>This is copied from the dataflow runner code.
   *
   * @param <T> the type of elements to concatenate.
   */
  private static class Concatenate<T> extends Combine.CombineFn<T, List<T>, List<T>> {
    @Override
    public List<T> createAccumulator() {
      return new ArrayList<>();
    }

    @Override
    public List<T> addInput(List<T> accumulator, T input) {
      accumulator.add(input);
      return accumulator;
    }

    @Override
    public List<T> mergeAccumulators(Iterable<List<T>> accumulators) {
      List<T> result = createAccumulator();
      for (List<T> accumulator : accumulators) {
        result.addAll(accumulator);
      }
      return result;
    }

    @Override
    public List<T> extractOutput(List<T> accumulator) {
      return accumulator;
    }

    @Override
    public Coder<List<T>> getAccumulatorCoder(CoderRegistry registry, Coder<T> inputCoder) {
      return ListCoder.of(inputCoder);
    }

    @Override
    public Coder<List<T>> getDefaultOutputCoder(CoderRegistry registry, Coder<T> inputCoder) {
      return ListCoder.of(inputCoder);
    }
  }

  private static void urnNotFound(
      PTransformNode transform, RunnerApi.Pipeline pipeline, BatchTranslationContext context) {
    throw new IllegalArgumentException(
        String.format(
            "Unknown type of URN %s for PTransform with id %s.",
            transform.getTransform().getSpec().getUrn(), transform.getId()));
  }

  private static void pruneOutput(
      DataSet<RawUnionValue> taggedDataset,
      BatchTranslationContext context,
      int unionTag,
      Coder<WindowedValue<?>> outputCoder,
      String transformName,
      String collectionId) {
    TypeInformation<WindowedValue<?>> outputType = new CoderTypeInformation<>(outputCoder);
    FlinkExecutableStagePruningFunction pruningFunction =
        new FlinkExecutableStagePruningFunction(unionTag);
    FlatMapOperator<RawUnionValue, WindowedValue<?>> pruningOperator =
        new FlatMapOperator<>(
            taggedDataset,
            outputType,
            pruningFunction,
            String.format("%s/out.%d", transformName, unionTag));
    context.addDataSet(collectionId, pruningOperator);
  }
}