Refactor LateDataDroppingDoFnRunner

This moves the LateDataDroppingDoFnRunner into util, removes references
to worker code from util, and genericizes the general
GroupAlsoByWindowsDoFn.

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

Author: tgroh

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

@@ -23,6 +23,8 @@
 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.ComposedKeyedWorkItem;
+import com.google.cloud.dataflow.sdk.util.KeyedWorkItem;
 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;
@@ -79,7 +81,7 @@ public static <K, ElemT> KeyedWorkItem<K, ElemT> windmillWorkItem(
    */
   public static <K, ElemT> KeyedWorkItem<K, ElemT> elementsWorkItem(
       K key, Iterable<WindowedValue<ElemT>> elementsIterable) {
-    return new ComposedKeyedWorkItem<>(key, Collections.<TimerData>emptyList(), elementsIterable);
+    return ComposedKeyedWorkItem.create(key, Collections.<TimerData>emptyList(), elementsIterable);
   }
 
   /**
@@ -90,7 +92,7 @@ public static <K, ElemT> KeyedWorkItem<K, ElemT> elementsWorkItem(
    */
   public static <K, ElemT> KeyedWorkItem<K, ElemT> timersWorkItem(
       K key, Iterable<TimerData> timersIterable) {
-    return new ComposedKeyedWorkItem<>(
+    return ComposedKeyedWorkItem.create(
         key, timersIterable, Collections.<WindowedValue<ElemT>>emptyList());
   }
 
@@ -103,8 +105,7 @@ public static <K, ElemT> KeyedWorkItem<K, ElemT> timersWorkItem(
    */
   public static <K, ElemT> KeyedWorkItem<K, ElemT> workItem(
       K key, Iterable<TimerData> timersIterable, Iterable<WindowedValue<ElemT>> elementsIterable) {
-    return new ComposedKeyedWorkItem<>(
-        key, timersIterable, elementsIterable);
+    return ComposedKeyedWorkItem.create(key, timersIterable, elementsIterable);
   }
 
   private static class WindmillKeyedWorkItem<K, ElemT> implements KeyedWorkItem<K, ElemT> {
@@ -217,33 +218,6 @@ public int hashCode() {
     }
   }
 
-  private static class ComposedKeyedWorkItem<K, ElemT> implements KeyedWorkItem<K, ElemT> {
-    private final K key;
-    private final Iterable<TimerData> timersIterable;
-    private final Iterable<WindowedValue<ElemT>> elementsIterable;
-
-    ComposedKeyedWorkItem(K key, Iterable<TimerData> timersIterable,
-        Iterable<WindowedValue<ElemT>> elementsIterable) {
-      this.key = key;
-      this.timersIterable = timersIterable;
-      this.elementsIterable = elementsIterable;
-    }
-    @Override
-    public K key() {
-      return key;
-    }
-
-    @Override
-    public Iterable<TimerData> timersIterable() {
-      return timersIterable;
-    }
-
-    @Override
-    public Iterable<WindowedValue<ElemT>> elementsIterable() {
-      return elementsIterable;
-    }
-  }
-
   /**
    * Coder that forwards {@code ByteSizeObserver} calls to an underlying element coder.
    * {@code TimerOrElement} objects never need to be encoded, so this class does not

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

@@ -21,7 +21,6 @@
 import com.google.cloud.dataflow.sdk.transforms.DoFn;
 import com.google.cloud.dataflow.sdk.util.DoFnInfo;
 import com.google.cloud.dataflow.sdk.util.DoFnRunner;
-import com.google.cloud.dataflow.sdk.util.DoFnRunner.RequiresLateDataDropping;
 import com.google.cloud.dataflow.sdk.util.DoFnRunners;
 import com.google.cloud.dataflow.sdk.util.DoFnRunners.OutputManager;
 import com.google.cloud.dataflow.sdk.util.ExecutionContext;
@@ -179,21 +178,8 @@ public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {
       }
     };
 
-    if (doFnInfo.getDoFn() instanceof StreamingGroupAlsoByWindowsDoFn
-        && doFnInfo.getDoFn() instanceof RequiresLateDataDropping) {
-      @SuppressWarnings({"unchecked", "rawtypes"})
-      DoFnRunner<Object, Object> objectFnRunner = DoFnRunners.lateDataDroppingRunner(
-          options,
-          (DoFnInfo) doFnInfo,
-          sideInputReader,
-          outputManager,
-          (TupleTag) mainOutputTag,
-          sideOutputTags,
-          stepContext,
-          addCounterMutator);
-      fnRunner = objectFnRunner;
-    } else if (hasStreamingSideInput) {
-      fnRunner = DoFnRunners.streamingSideInputRunner(
+    if (hasStreamingSideInput) {
+      fnRunner = StreamingDoFnRunners.streamingSideInputRunner(
           options,
           doFnInfo,
           sideInputReader,
@@ -203,7 +189,7 @@ public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {
           stepContext,
           addCounterMutator);
     } else {
-      fnRunner = DoFnRunners.simpleRunner(
+      fnRunner = DoFnRunners.createDefault(
           options,
           doFnInfo.getDoFn(),
           sideInputReader,

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

@@ -0,0 +1,65 @@
+/*
+ * Copyright (C) 2015 Google Inc.
+ *
+ * Licensed 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 com.google.cloud.dataflow.sdk.runners.worker;
+
+import com.google.cloud.dataflow.sdk.options.PipelineOptions;
+import com.google.cloud.dataflow.sdk.util.DoFnInfo;
+import com.google.cloud.dataflow.sdk.util.DoFnRunner;
+import com.google.cloud.dataflow.sdk.util.DoFnRunners.OutputManager;
+import com.google.cloud.dataflow.sdk.util.ExecutionContext.StepContext;
+import com.google.cloud.dataflow.sdk.util.SideInputReader;
+import com.google.cloud.dataflow.sdk.util.common.CounterSet;
+import com.google.cloud.dataflow.sdk.values.TupleTag;
+
+import java.util.List;
+
+/**
+ * Utility methods for creating {@link DoFnRunner} instances used by streaming Dataflow.
+ */
+public final class StreamingDoFnRunners {
+  private StreamingDoFnRunners() {
+    // Do not instantiate
+  }
+
+  /**
+   * Returns an implementation of {@link DoFnRunner} that handles streaming side inputs.
+   *
+   * <p>It blocks and caches input elements if their side inputs are not ready.
+   */
+  public static <InputT, OutputT> DoFnRunner<InputT, OutputT> streamingSideInputRunner(
+      PipelineOptions options,
+      DoFnInfo<InputT, OutputT> doFnInfo,
+      SideInputReader sideInputReader,
+      OutputManager outputManager,
+      TupleTag<OutputT> mainOutputTag,
+      List<TupleTag<?>> sideOutputTags,
+      StepContext stepContext,
+      CounterSet.AddCounterMutator addCounterMutator) {
+    return new StreamingSideInputDoFnRunner<>(
+        options,
+        doFnInfo,
+        sideInputReader,
+        outputManager,
+        mainOutputTag,
+        sideOutputTags,
+        stepContext,
+        addCounterMutator,
+        doFnInfo.getWindowingStrategy());
+  }
+
+
+}
+

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

@@ -17,18 +17,13 @@
 package com.google.cloud.dataflow.sdk.runners.worker;
 
 import com.google.cloud.dataflow.sdk.coders.Coder;
-import com.google.cloud.dataflow.sdk.transforms.Aggregator;
 import com.google.cloud.dataflow.sdk.transforms.DoFn;
-import com.google.cloud.dataflow.sdk.transforms.Sum;
 import com.google.cloud.dataflow.sdk.transforms.windowing.BoundedWindow;
 import com.google.cloud.dataflow.sdk.util.AppliedCombineFn;
-import com.google.cloud.dataflow.sdk.util.DoFnRunner.RequiresLateDataDropping;
-import com.google.cloud.dataflow.sdk.util.GroupAlsoByWindowsDoFn;
-import com.google.cloud.dataflow.sdk.util.ReduceFnRunner;
+import com.google.cloud.dataflow.sdk.util.GroupAlsoByWindowViaWindowSetDoFn;
+import com.google.cloud.dataflow.sdk.util.KeyedWorkItem;
 import com.google.cloud.dataflow.sdk.util.SystemDoFnInternal;
 import com.google.cloud.dataflow.sdk.util.SystemReduceFn;
-import com.google.cloud.dataflow.sdk.util.TimerInternals;
-import com.google.cloud.dataflow.sdk.util.TimerInternals.TimerData;
 import com.google.cloud.dataflow.sdk.util.WindowingStrategy;
 import com.google.cloud.dataflow.sdk.values.KV;
 import com.google.common.base.Preconditions;
@@ -44,64 +39,28 @@
 @SystemDoFnInternal
 public abstract class StreamingGroupAlsoByWindowsDoFn<K, InputT, OutputT, W extends BoundedWindow>
     extends DoFn<KeyedWorkItem<K, InputT>, KV<K, OutputT>> {
-  protected final Aggregator<Long, Long> droppedDueToClosedWindow = createAggregator(
-      GroupAlsoByWindowsDoFn.DROPPED_DUE_TO_CLOSED_WINDOW_COUNTER, new Sum.SumLongFn());
-  protected final Aggregator<Long, Long> droppedDueToLateness = createAggregator(
-      GroupAlsoByWindowsDoFn.DROPPED_DUE_TO_LATENESS_COUNTER, new Sum.SumLongFn());
-
   public static <K, InputT, AccumT, OutputT, W extends BoundedWindow>
-      StreamingGroupAlsoByWindowsDoFn<K, InputT, OutputT, W> create(
+      DoFn<KeyedWorkItem<K, InputT>, KV<K, OutputT>> create(
           final WindowingStrategy<?, W> windowingStrategy,
           final AppliedCombineFn<K, InputT, AccumT, OutputT> combineFn,
           final Coder<K> keyCoder) {
     Preconditions.checkNotNull(combineFn);
-    return new StreamingGABWViaWindowSetDoFn<>(windowingStrategy,
-        SystemReduceFn.<K, InputT, AccumT, OutputT, W>combining(keyCoder, combineFn));
+    DoFn<KeyedWorkItem<K, InputT>, KV<K, OutputT>> fn =
+        GroupAlsoByWindowViaWindowSetDoFn.create(
+            windowingStrategy,
+            SystemReduceFn.<K, InputT, AccumT, OutputT, W>combining(keyCoder, combineFn));
+    return fn;
   }
 
   public static <K, V, W extends BoundedWindow>
-  DoFn<KeyedWorkItem<K, V>, KV<K, Iterable<V>>> createForIterable(
-      final WindowingStrategy<?, W> windowingStrategy,
-      final Coder<V> inputCoder) {
+      DoFn<KeyedWorkItem<K, V>, KV<K, Iterable<V>>> createForIterable(
+          final WindowingStrategy<?, W> windowingStrategy, final Coder<V> inputCoder) {
     // If the windowing strategy indicates we're doing a reshuffle, use the special-path.
     if (StreamingGroupAlsoByWindowsReshuffleDoFn.isReshuffle(windowingStrategy)) {
       return new StreamingGroupAlsoByWindowsReshuffleDoFn<>();
     } else {
-      return new StreamingGABWViaWindowSetDoFn<>(
+      return GroupAlsoByWindowViaWindowSetDoFn.create(
           windowingStrategy, SystemReduceFn.<K, V, W>buffering(inputCoder));
     }
   }
-
-  private static class StreamingGABWViaWindowSetDoFn<K, InputT, OutputT, W extends BoundedWindow>
-  extends StreamingGroupAlsoByWindowsDoFn<K, InputT, OutputT, W>
-  implements RequiresLateDataDropping {
-    private final WindowingStrategy<Object, W> windowingStrategy;
-    private SystemReduceFn.Factory<K, InputT, OutputT, W> reduceFnFactory;
-
-    public StreamingGABWViaWindowSetDoFn(WindowingStrategy<?, W> windowingStrategy,
-        SystemReduceFn.Factory<K, InputT, OutputT, W> reduceFnFactory) {
-      @SuppressWarnings("unchecked")
-      WindowingStrategy<Object, W> noWildcard = (WindowingStrategy<Object, W>) windowingStrategy;
-      this.windowingStrategy = noWildcard;
-      this.reduceFnFactory = reduceFnFactory;
-    }
-
-    @Override
-    public void processElement(ProcessContext c) throws Exception {
-      KeyedWorkItem<K, InputT> element = c.element();
-
-      @SuppressWarnings("unchecked")
-      K key = c.element().key();
-      TimerInternals timerInternals = c.windowingInternals().timerInternals();
-      ReduceFnRunner<K, InputT, OutputT, W> runner = new ReduceFnRunner<>(
-            key, windowingStrategy, timerInternals, c.windowingInternals(),
-            droppedDueToClosedWindow, reduceFnFactory.create(key));
-
-      for (TimerData timer : element.timersIterable()) {
-        runner.onTimer(timer);
-      }
-      runner.processElements(element.elementsIterable());
-      runner.persist();
-    }
-  }
 }

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

@@ -16,6 +16,7 @@
 package com.google.cloud.dataflow.sdk.runners.worker;
 
 import com.google.cloud.dataflow.sdk.transforms.DoFn;
+import com.google.cloud.dataflow.sdk.util.KeyedWorkItem;
 import com.google.cloud.dataflow.sdk.util.ReshuffleTrigger;
 import com.google.cloud.dataflow.sdk.util.SystemDoFnInternal;
 import com.google.cloud.dataflow.sdk.util.WindowedValue;

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

@@ -24,6 +24,7 @@
 import com.google.cloud.dataflow.sdk.transforms.windowing.BoundedWindow;
 import com.google.cloud.dataflow.sdk.util.CloudObject;
 import com.google.cloud.dataflow.sdk.util.ExecutionContext;
+import com.google.cloud.dataflow.sdk.util.KeyedWorkItem;
 import com.google.cloud.dataflow.sdk.util.WindowedValue;
 import com.google.cloud.dataflow.sdk.util.WindowedValue.FullWindowedValueCoder;
 import com.google.cloud.dataflow.sdk.util.common.CounterSet;
@@ -99,7 +100,7 @@ public NativeReaderIterator<WindowedValue<KeyedWorkItem<K, T>>> iterator() throw
     final WorkItem workItem = context.getWork();
     final WindowedValue<KeyedWorkItem<K, T>> value =
         WindowedValue.valueInEmptyWindows(
-            KeyedWorkItems.windmillWorkItem(
+            KeyedWorkItems.<K, T>windmillWorkItem(
                 key, workItem, windowCoder, windowsCoder, kvCoder.getValueCoder()));
 
     return new NativeReaderIterator<WindowedValue<KeyedWorkItem<K, T>>>() {

sdk/src/main/java/com/google/cloud/dataflow/sdk/transforms/DoFnTester.java

@@ -481,7 +481,7 @@ void initializeState() {
       runnerSideInputs = runnerSideInputs.and(entry.getKey().getTagInternal(), entry.getValue());
     }
     outputManager = new DoFnRunnerBase.ListOutputManager();
-    fnRunner = DoFnRunners.simpleRunner(
+    fnRunner = DoFnRunners.createDefault(
         options,
         fn,
         DirectSideInputReader.of(runnerSideInputs),

sdk/src/main/java/com/google/cloud/dataflow/sdk/transforms/ParDo.java

@@ -1144,7 +1144,7 @@ private static <InputT, OutputT, ActualInputT extends InputT> void evaluateHelpe
             outputs);
 
     DoFnRunner<InputT, OutputT> fnRunner =
-        DoFnRunners.simpleRunner(
+        DoFnRunners.createDefault(
             context.getPipelineOptions(),
             fn,
             sideInputReader,

sdk/src/main/java/com/google/cloud/dataflow/sdk/util/ComposedKeyedWorkItem.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2015 Google Inc.
+ *
+ * Licensed 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 com.google.cloud.dataflow.sdk.util;
+
+import com.google.cloud.dataflow.sdk.util.TimerInternals.TimerData;
+
+/**
+ * A {@link KeyedWorkItem} composed of an underlying key, {@link TimerData} iterable, and element
+ * iterable.
+ */
+public class ComposedKeyedWorkItem<K, ElemT> implements KeyedWorkItem<K, ElemT> {
+
+  private final K key;
+  private final Iterable<TimerData> timers;
+  private final Iterable<WindowedValue<ElemT>> elements;
+
+  public static <K, ElemT> ComposedKeyedWorkItem<K, ElemT> create(
+      K key, Iterable<TimerData> timers, Iterable<WindowedValue<ElemT>> elements) {
+    return new ComposedKeyedWorkItem<K, ElemT>(key, timers, elements);
+  }
+
+  private ComposedKeyedWorkItem(
+      K key, Iterable<TimerData> timers, Iterable<WindowedValue<ElemT>> elements) {
+    this.key = key;
+    this.timers = timers;
+    this.elements = elements;
+  }
+
+  @Override
+  public K key() {
+    return key;
+  }
+
+  @Override
+  public Iterable<TimerData> timersIterable() {
+    return timers;
+  }
+
+  @Override
+  public Iterable<WindowedValue<ElemT>> elementsIterable() {
+    return elements;
+  }
+}