Create StreamingGroupAlsoByWindowsDoFnRunner

Move the input elements lateness check to DoFnRunner.

----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=112982861

Author: peihe

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/GroupAlsoByWindowsParDoFn.java

@@ -65,14 +65,15 @@ class GroupAlsoByWindowsParDoFn extends ParDoFnBase {
   static GroupAlsoByWindowsParDoFn of(
       PipelineOptions options,
       DoFn<?, ?> groupAlsoByWindowsDoFn,
+      WindowingStrategy<?, ?> windowingStrategy,
       String stepName,
       String transformName,
-      DataflowExecutionContext executionContext,
+      DataflowExecutionContext<?> executionContext,
       CounterSet.AddCounterMutator addCounterMutator,
       StateSampler stateSampler)
       throws Exception {
-    return new GroupAlsoByWindowsParDoFn(options, groupAlsoByWindowsDoFn, stepName, transformName,
-        executionContext, addCounterMutator, stateSampler);
+    return new GroupAlsoByWindowsParDoFn(options, groupAlsoByWindowsDoFn, windowingStrategy,
+        stepName, transformName, executionContext, addCounterMutator, stateSampler);
   }
 
   /**
@@ -160,6 +161,7 @@ public ParDoFn create(
       return GroupAlsoByWindowsParDoFn.of(
           options,
           groupAlsoByWindowsDoFn,
+          windowingStrategy,
           stepName,
           transformName,
           executionContext,
@@ -170,7 +172,7 @@ public ParDoFn create(
 
   @Override
   protected DoFnInfo<?, ?> getDoFnInfo() {
-    return new DoFnInfo<>(groupAlsoByWindowsDoFn, null);
+    return new DoFnInfo<>(groupAlsoByWindowsDoFn, windowingStrategy);
   }
 
   @SuppressWarnings({"unchecked", "rawtypes"})
@@ -262,13 +264,15 @@ public Coder<List<AccumT>> getAccumulatorCoder(CoderRegistry registry, Coder<K>
   }
 
   private final DoFn<?, ?> groupAlsoByWindowsDoFn;
+  private final WindowingStrategy<?, ?> windowingStrategy;
 
   private GroupAlsoByWindowsParDoFn(
       PipelineOptions options,
       DoFn<?, ?> groupAlsoByWindowsDoFn,
+      WindowingStrategy<?, ?> windowingStrategy,
       String stepName,
       String transformName,
-      DataflowExecutionContext executionContext,
+      DataflowExecutionContext<?> executionContext,
       CounterSet.AddCounterMutator addCounterMutator,
       StateSampler stateSampler) {
     super(
@@ -281,5 +285,6 @@ private GroupAlsoByWindowsParDoFn(
         addCounterMutator,
         stateSampler);
     this.groupAlsoByWindowsDoFn = groupAlsoByWindowsDoFn;
+    this.windowingStrategy = windowingStrategy;
   }
 }

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/KeyedWorkItem.java

@@ -13,228 +13,33 @@
  * License for the specific language governing permissions and limitations under
  * the License.
  */
-
 package com.google.cloud.dataflow.sdk.runners.worker;
 
-import com.google.cloud.dataflow.sdk.coders.Coder;
-import com.google.cloud.dataflow.sdk.coders.StandardCoder;
-import com.google.cloud.dataflow.sdk.runners.worker.windmill.Windmill;
-import com.google.cloud.dataflow.sdk.runners.worker.windmill.Windmill.InputMessageBundle;
-import com.google.cloud.dataflow.sdk.runners.worker.windmill.Windmill.Message;
-import com.google.cloud.dataflow.sdk.runners.worker.windmill.Windmill.Timer;
-import com.google.cloud.dataflow.sdk.transforms.windowing.BoundedWindow;
-import com.google.cloud.dataflow.sdk.transforms.windowing.PaneInfo;
-import com.google.cloud.dataflow.sdk.util.PropertyNames;
-import com.google.cloud.dataflow.sdk.util.TimeDomain;
 import com.google.cloud.dataflow.sdk.util.TimerInternals.TimerData;
 import com.google.cloud.dataflow.sdk.util.WindowedValue;
-import com.google.cloud.dataflow.sdk.util.common.ElementByteSizeObserver;
-import com.google.cloud.dataflow.sdk.util.state.StateNamespace;
-import com.google.cloud.dataflow.sdk.util.state.StateNamespaces;
-import com.google.common.base.Function;
-import com.google.common.base.Predicate;
-import com.google.common.base.Predicates;
-import com.google.common.base.Throwables;
-import com.google.common.collect.FluentIterable;
-
-import com.fasterxml.jackson.annotation.JsonCreator;
-import com.fasterxml.jackson.annotation.JsonProperty;
-
-import org.joda.time.Instant;
-
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.List;
-import java.util.Objects;
-import java.util.concurrent.TimeUnit;
 
 /**
- * Wrapper around a {@link Windmill.WorkItem} which contains all the timers and elements associated
- * with a specific work item.
+ * Interface that contains all the timers and elements associated with a specific work item.
  *
  * <p>Used as the input type of {@link StreamingGroupAlsoByWindowsDoFn}.
  *
+ * @param <K> the key type
  * @param <ElemT> the element type
  */
-public class KeyedWorkItem<ElemT> {
-
-  private static final Predicate<Timer> IS_WATERMARK = new Predicate<Timer>() {
-    @Override
-    public boolean apply(Timer input) {
-      return input.getType() == Timer.Type.WATERMARK;
-    }
-  };
-
-  public static <ElemT> KeyedWorkItem<ElemT> workItem(
-      Object key,
-      Windmill.WorkItem workItem,
-      Coder<? extends BoundedWindow> windowCoder,
-      Coder<Collection<? extends BoundedWindow>> windowsCoder,
-      Coder<ElemT> valueCoder) {
-    return new KeyedWorkItem<>(key, workItem, windowCoder, windowsCoder, valueCoder);
-  }
-
-  public Object key() {
-    return key;
-  }
-
-  public Iterable<TimerData> timersIterable() {
-    FluentIterable<Timer> allTimers = FluentIterable.from(workItem.getTimers().getTimersList());
-    FluentIterable<Timer> eventTimers = allTimers.filter(IS_WATERMARK);
-    FluentIterable<Timer> nonEventTimers = allTimers.filter(Predicates.not(IS_WATERMARK));
-    return eventTimers.append(nonEventTimers).transform(new Function<Timer, TimerData>() {
-      private TimeDomain getTimeDomain(Windmill.Timer.Type type) {
-        switch (type) {
-          case REALTIME:
-            return TimeDomain.PROCESSING_TIME;
-          case DEPENDENT_REALTIME:
-            return TimeDomain.SYNCHRONIZED_PROCESSING_TIME;
-          case WATERMARK:
-            return TimeDomain.EVENT_TIME;
-          default:
-            throw new IllegalArgumentException("Unsupported timer type " + type);
-        }
-      }
-
-      @Override
-      public TimerData apply(Timer timer) {
-        String tag = timer.getTag().toStringUtf8();
-        String namespaceString = tag.substring(0, tag.indexOf('+'));
-        StateNamespace namespace = StateNamespaces.fromString(namespaceString, windowCoder);
-
-        Instant timestamp = new Instant(TimeUnit.MICROSECONDS.toMillis(timer.getTimestamp()));
-        return TimerData.of(namespace, timestamp, getTimeDomain(timer.getType()));
-      }
-    });
-  }
-
-  public Iterable<WindowedValue<ElemT>> elementsIterable() {
-    return FluentIterable.from(workItem.getMessageBundlesList())
-        .transformAndConcat(new Function<InputMessageBundle, Iterable<Message>>() {
-          @Override
-          public Iterable<Message> apply(InputMessageBundle input) {
-            return input.getMessagesList();
-          }
-        })
-        .transform(new Function<Message, WindowedValue<ElemT>>() {
-          @Override
-          public WindowedValue<ElemT> apply(Message message) {
-            try {
-              Instant timestamp = new Instant(
-                  TimeUnit.MICROSECONDS.toMillis(message.getTimestamp()));
-              Collection<? extends BoundedWindow> windows =
-                  WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata());
-              PaneInfo pane = WindmillSink.decodeMetadataPane(message.getMetadata());
-
-              InputStream inputStream = message.getData().newInput();
-              ElemT value = valueCoder.decode(inputStream, Coder.Context.OUTER);
-              return WindowedValue.of(value, timestamp, windows, pane);
-            } catch (IOException e) {
-              throw Throwables.propagate(e);
-            }
-          }
-        });
-  }
-
-  @Override
-  public boolean equals(Object other) {
-    if (!(other instanceof KeyedWorkItem)) {
-      return false;
-    }
-
-    KeyedWorkItem<?> that = (KeyedWorkItem<?>) other;
-    return Objects.equals(this.key, that.key)
-        && Objects.equals(this.workItem, that.workItem);
-  }
-
-  @Override
-  public int hashCode() {
-    return Objects.hash(key, workItem);
-  }
+public interface KeyedWorkItem<K, ElemT> {
 
   /**
-   * Coder that forwards {@code ByteSizeObserver} calls to an underlying element coder.
-   * {@code TimerOrElement} objects never need to be encoded, so this class does not
-   * support the {@code encode} and {@code decode} methods.
+   * Returns the key.
    */
-  public static class KeyedWorkItemCoder<T> extends StandardCoder<KeyedWorkItem<T>> {
-    final Coder<T> elemCoder;
-
-    /**
-     * Creates a new {@code TimerOrElement.Coder} that wraps the given {@link Coder}.
-     */
-    public static <T> KeyedWorkItemCoder<T> of(Coder<T> elemCoder) {
-      return new KeyedWorkItemCoder<>(elemCoder);
-    }
+  public K key();
 
-    @JsonCreator
-    public static KeyedWorkItemCoder<?> of(
-            @JsonProperty(PropertyNames.COMPONENT_ENCODINGS)
-            List<Object> components) {
-      return of((Coder<?>) components.get(0));
-    }
-
-    @Override
-    public void encode(KeyedWorkItem<T> value, OutputStream outStream, Context context) {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public KeyedWorkItem<T> decode(InputStream inStream, Context context) {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public boolean isRegisterByteSizeObserverCheap(KeyedWorkItem<T> value, Context context) {
-      return true;
-    }
-
-    @Override
-    public void registerByteSizeObserver(
-        KeyedWorkItem<T> value, ElementByteSizeObserver observer, Context context)
-        throws Exception {
-      observer.update((long) value.workItem.getSerializedSize());
-    }
-
-    @Override
-    public void verifyDeterministic() throws NonDeterministicException {}
-
-    @Override
-    public List<? extends Coder<?>> getCoderArguments() {
-      return Arrays.asList(elemCoder);
-    }
-
-    public Coder<T> getElementCoder() {
-      return elemCoder;
-    }
-
-    protected KeyedWorkItemCoder(Coder<T> elemCoder) {
-      this.elemCoder = elemCoder;
-    }
-  }
-
-  //////////////////////////////////////////////////////////////////////////////
-
-  private final Windmill.WorkItem workItem;
-  private final Object key;
-
-  private final transient Coder<? extends BoundedWindow> windowCoder;
-  private final transient Coder<Collection<? extends BoundedWindow>> windowsCoder;
-  private final transient Coder<ElemT> valueCoder;
+  /**
+   * Returns the timers iterable.
+   */
+  public Iterable<TimerData> timersIterable();
 
-  KeyedWorkItem(
-      Object key,
-      Windmill.WorkItem workItem,
-      Coder<? extends BoundedWindow> windowCoder,
-      Coder<Collection<? extends BoundedWindow>> windowsCoder,
-      Coder<ElemT> valueCoder) {
-    this.key = key;
-    this.workItem = workItem;
-    this.windowCoder = windowCoder;
-    this.windowsCoder = windowsCoder;
-    this.valueCoder = valueCoder;
-  }
+  /**
+   * Returns the elements iterable.
+   */
+  public Iterable<WindowedValue<ElemT>> elementsIterable();
 }