[BEAM-22] Implement InProcessEvaluationContext

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/inprocess/BoundedReadEvaluatorFactory.java

@@ -18,7 +18,6 @@
 import com.google.cloud.dataflow.sdk.io.Read.Bounded;
 import com.google.cloud.dataflow.sdk.io.Source.Reader;
 import com.google.cloud.dataflow.sdk.runners.inprocess.InProcessPipelineRunner.CommittedBundle;
-import com.google.cloud.dataflow.sdk.runners.inprocess.InProcessPipelineRunner.InProcessEvaluationContext;
 import com.google.cloud.dataflow.sdk.runners.inprocess.InProcessPipelineRunner.UncommittedBundle;
 import com.google.cloud.dataflow.sdk.transforms.AppliedPTransform;
 import com.google.cloud.dataflow.sdk.transforms.PTransform;

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/inprocess/EvaluatorKey.java

@@ -15,7 +15,6 @@
  */
 package com.google.cloud.dataflow.sdk.runners.inprocess;
 
-import com.google.cloud.dataflow.sdk.runners.inprocess.InProcessPipelineRunner.InProcessEvaluationContext;
 import com.google.cloud.dataflow.sdk.transforms.AppliedPTransform;
 
 import java.util.Objects;

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/inprocess/FlattenEvaluatorFactory.java

@@ -16,7 +16,6 @@
 package com.google.cloud.dataflow.sdk.runners.inprocess;
 
 import com.google.cloud.dataflow.sdk.runners.inprocess.InProcessPipelineRunner.CommittedBundle;
-import com.google.cloud.dataflow.sdk.runners.inprocess.InProcessPipelineRunner.InProcessEvaluationContext;
 import com.google.cloud.dataflow.sdk.runners.inprocess.InProcessPipelineRunner.UncommittedBundle;
 import com.google.cloud.dataflow.sdk.transforms.AppliedPTransform;
 import com.google.cloud.dataflow.sdk.transforms.Flatten;

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/inprocess/GroupByKeyEvaluatorFactory.java

@@ -22,7 +22,6 @@
 import com.google.cloud.dataflow.sdk.coders.IterableCoder;
 import com.google.cloud.dataflow.sdk.coders.KvCoder;
 import com.google.cloud.dataflow.sdk.runners.inprocess.InProcessPipelineRunner.CommittedBundle;
-import com.google.cloud.dataflow.sdk.runners.inprocess.InProcessPipelineRunner.InProcessEvaluationContext;
 import com.google.cloud.dataflow.sdk.runners.inprocess.InProcessPipelineRunner.UncommittedBundle;
 import com.google.cloud.dataflow.sdk.runners.inprocess.StepTransformResult.Builder;
 import com.google.cloud.dataflow.sdk.transforms.AppliedPTransform;

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/inprocess/InMemoryWatermarkManager.java

@@ -1209,8 +1209,11 @@ public TimerUpdateBuilder deletedTimer(TimerData deletedTimer) {
        * and deletedTimers.
        */
       public TimerUpdate build() {
-        return new TimerUpdate(key, ImmutableSet.copyOf(completedTimers),
-            ImmutableSet.copyOf(setTimers), ImmutableSet.copyOf(deletedTimers));
+        return new TimerUpdate(
+            key,
+            ImmutableSet.copyOf(completedTimers),
+            ImmutableSet.copyOf(setTimers),
+            ImmutableSet.copyOf(deletedTimers));
       }
     }
 
@@ -1245,6 +1248,13 @@ Iterable<? extends TimerData> getDeletedTimers() {
       return deletedTimers;
     }
 
+    /**
+     * Returns a {@link TimerUpdate} that is like this one, but with the specified completed timers.
+     */
+    public TimerUpdate withCompletedTimers(Iterable<TimerData> completedTimers) {
+      return new TimerUpdate(this.key, completedTimers, setTimers, deletedTimers);
+    }
+
     @Override
     public int hashCode() {
       return Objects.hash(key, completedTimers, setTimers, deletedTimers);

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/inprocess/InProcessBundle.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2015 Google Inc.
+ * Copyright (C) 2016 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
@@ -22,7 +22,6 @@
 import com.google.cloud.dataflow.sdk.util.WindowedValue;
 import com.google.cloud.dataflow.sdk.values.PCollection;
 import com.google.common.base.MoreObjects;
-import com.google.common.base.MoreObjects.ToStringHelper;
 import com.google.common.collect.ImmutableList;
 
 import org.joda.time.Instant;
@@ -64,6 +63,11 @@ private InProcessBundle(PCollection<T> pcollection, boolean keyed, Object key) {
     this.elements = ImmutableList.builder();
   }
 
+  @Override
+  public PCollection<T> getPCollection() {
+    return pcollection;
+  }
+
   @Override
   public InProcessBundle<T> add(WindowedValue<T> element) {
     checkState(!committed, "Can't add element %s to committed bundle %s", element, this);
@@ -105,12 +109,12 @@ public Instant getSynchronizedProcessingOutputWatermark() {
 
       @Override
       public String toString() {
-        ToStringHelper toStringHelper =
-            MoreObjects.toStringHelper(this).add("pcollection", pcollection);
-        if (keyed) {
-          toStringHelper = toStringHelper.add("key", key);
-        }
-        return toStringHelper.add("elements", elements).toString();
+        return MoreObjects.toStringHelper(this)
+            .omitNullValues()
+            .add("pcollection", pcollection)
+            .add("key", key)
+            .add("elements", committedElements)
+            .toString();
       }
     };
   }