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

import com.google.common.collect.ImmutableMap;
import java.util.Collections;
import org.apache.beam.runners.core.GroupAlsoByWindowViaWindowSetDoFn;
import org.apache.beam.runners.core.GroupByKeyViaGroupByKeyOnly;
import org.apache.beam.runners.core.GroupByKeyViaGroupByKeyOnly.GroupAlsoByWindow;
import org.apache.beam.runners.core.GroupByKeyViaGroupByKeyOnly.GroupByKeyOnly;
import org.apache.beam.runners.core.SystemReduceFn;
import org.apache.beam.runners.direct.DirectExecutionContext.DirectStepContext;
import org.apache.beam.runners.direct.DirectGroupByKey.DirectGroupAlsoByWindow;
import org.apache.beam.runners.direct.DirectRunner.CommittedBundle;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.transforms.AppliedPTransform;
import org.apache.beam.sdk.transforms.OldDoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.KeyedWorkItem;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.util.WindowingStrategy;
import org.apache.beam.sdk.util.state.StateInternals;
import org.apache.beam.sdk.util.state.StateInternalsFactory;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;

/**
 * The {@link DirectRunner} {@link TransformEvaluatorFactory} for the
 * {@link GroupByKeyOnly} {@link PTransform}.
 */
class GroupAlsoByWindowEvaluatorFactory implements TransformEvaluatorFactory {
  private final EvaluationContext evaluationContext;

  GroupAlsoByWindowEvaluatorFactory(EvaluationContext evaluationContext) {
    this.evaluationContext = evaluationContext;
  }

  @Override
  public <InputT> TransformEvaluator<InputT> forApplication(
      AppliedPTransform<?, ?, ?> application,
      CommittedBundle<?> inputBundle) {
    @SuppressWarnings({"cast", "unchecked", "rawtypes"})
    TransformEvaluator<InputT> evaluator =
        createEvaluator(
            (AppliedPTransform) application, (CommittedBundle) inputBundle);
    return evaluator;
  }

  @Override
  public void cleanup() {}

  private <K, V> TransformEvaluator<KeyedWorkItem<K, V>> createEvaluator(
      AppliedPTransform<
              PCollection<KeyedWorkItem<K, V>>,
              PCollection<KV<K, Iterable<V>>>,
              DirectGroupAlsoByWindow<K, V>>
          application,
      CommittedBundle<KeyedWorkItem<K, V>> inputBundle) {
    return new GroupAlsoByWindowEvaluator<>(
        evaluationContext, inputBundle, application);
  }

  /**
   * A transform evaluator for the pseudo-primitive {@link GroupAlsoByWindow}. Windowing is ignored;
   * all input should be in the global window since all output will be as well.
   *
   * @see GroupByKeyViaGroupByKeyOnly
   */
  private static class GroupAlsoByWindowEvaluator<K, V>
      implements TransformEvaluator<KeyedWorkItem<K, V>> {

    private final TransformEvaluator<KeyedWorkItem<K, V>> gabwParDoEvaluator;

    public GroupAlsoByWindowEvaluator(
        final EvaluationContext evaluationContext,
        CommittedBundle<KeyedWorkItem<K, V>> inputBundle,
        final AppliedPTransform<
                PCollection<KeyedWorkItem<K, V>>,
                PCollection<KV<K, Iterable<V>>>,
                DirectGroupAlsoByWindow<K, V>> application) {

      Coder<V> valueCoder =
          application.getTransform().getValueCoder(inputBundle.getPCollection().getCoder());

      @SuppressWarnings("unchecked")
      WindowingStrategy<?, BoundedWindow> windowingStrategy =
          (WindowingStrategy<?, BoundedWindow>) application.getTransform().getWindowingStrategy();

      DirectStepContext stepContext =
          evaluationContext
              .getExecutionContext(application, inputBundle.getKey())
              .getOrCreateStepContext(
                  evaluationContext.getStepName(application), application.getTransform().getName());

      StateInternals<K> stateInternals = (StateInternals<K>) stepContext.stateInternals();

      OldDoFn<KeyedWorkItem<K, V>, KV<K, Iterable<V>>> gabwDoFn =
          GroupAlsoByWindowViaWindowSetDoFn.create(
              windowingStrategy,
              new ConstantStateInternalsFactory<K>(stateInternals),
              SystemReduceFn.<K, V, BoundedWindow>buffering(valueCoder));

      TupleTag<KV<K, Iterable<V>>> mainOutputTag = new TupleTag<KV<K, Iterable<V>>>() {};

      // Not technically legit, as the application is not a ParDo
      this.gabwParDoEvaluator =
          ParDoEvaluator.create(
              evaluationContext,
              stepContext,
              inputBundle,
              application,
              gabwDoFn,
              Collections.<PCollectionView<?>>emptyList(),
              mainOutputTag,
              Collections.<TupleTag<?>>emptyList(),
              ImmutableMap.<TupleTag<?>, PCollection<?>>of(mainOutputTag, application.getOutput()));
    }

    @Override
    public void processElement(WindowedValue<KeyedWorkItem<K, V>> element) throws Exception {
      gabwParDoEvaluator.processElement(element);
    }

    @Override
    public TransformResult finishBundle() throws Exception {
      return gabwParDoEvaluator.finishBundle();
    }
  }

  private static final class ConstantStateInternalsFactory<K>
      implements StateInternalsFactory<K> {
    private final StateInternals<K> stateInternals;

    private ConstantStateInternalsFactory(StateInternals<K> stateInternals) {
      this.stateInternals = stateInternals;
    }

    @Override
    @SuppressWarnings("unchecked")
    public StateInternals<K> stateInternalsForKey(K key) {
      return stateInternals;
    }
  }
}

@@ -3782,16 +3782,103 @@
 K, V%3E%3E %7B
+%0A    private static final TupleTag%3CObject%3E MAIN_OUTPUT_TAG = new TupleTag%3CObject%3E() %7B%7D;
 %0A%0A    pr
@@ -5281,99 +5281,8 @@
 );%0A%0A
-      TupleTag%3CKV%3CK, Iterable%3CV%3E%3E%3E mainOutputTag = new TupleTag%3CKV%3CK, Iterable%3CV%3E%3E%3E() %7B%7D;%0A%0A
     
@@ -5619,29 +5619,31 @@
         
-mainOutputTag
+MAIN_OUTPUT_TAG
 ,%0A      
@@ -5751,21 +5751,42 @@
 %3Eof(
-mainOutputTag
+%0A                  MAIN_OUTPUT_TAG
 , ap

Use a Constant TupleTag in GroupAlsoByWindowEvaluator