[BEAM-22] Add CommittedResult

sdks/java/core/src/main/java/org/apache/beam/sdk/runners/inprocess/CommittedResult.java

@@ -0,0 +1,47 @@
+/*
+ * 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.sdk.runners.inprocess;
+
+import org.apache.beam.sdk.runners.inprocess.InProcessPipelineRunner.CommittedBundle;
+import org.apache.beam.sdk.transforms.AppliedPTransform;
+
+import com.google.auto.value.AutoValue;
+
+/**
+ * A {@link InProcessTransformResult} that has been committed.
+ */
+@AutoValue
+abstract class CommittedResult {
+  /**
+   * Returns the {@link AppliedPTransform} that produced this result.
+   */
+  public abstract AppliedPTransform<?, ?, ?> getTransform();
+
+  /**
+   * Returns the outputs produced by the transform.
+   */
+  public abstract Iterable<? extends CommittedBundle<?>> getOutputs();
+
+  public static CommittedResult create(
+      InProcessTransformResult original, Iterable<? extends CommittedBundle<?>> outputs) {
+    return new AutoValue_CommittedResult(original.getTransform(),
+        outputs);
+  }
+}

sdks/java/core/src/main/java/org/apache/beam/sdk/runners/inprocess/CompletionCallback.java

@@ -26,7 +26,7 @@ interface CompletionCallback {
   /**
    * Handle a successful result, returning the committed outputs of the result.
    */
-  Iterable<? extends CommittedBundle<?>> handleResult(
+  CommittedResult handleResult(
       CommittedBundle<?> inputBundle, InProcessTransformResult result);
 
   /**

sdks/java/core/src/main/java/org/apache/beam/sdk/runners/inprocess/ExecutorServiceParallelExecutor.java

@@ -216,14 +216,14 @@ public void awaitCompletion() throws Throwable {
    */
   private class DefaultCompletionCallback implements CompletionCallback {
     @Override
-    public Iterable<? extends CommittedBundle<?>> handleResult(
+    public CommittedResult handleResult(
         CommittedBundle<?> inputBundle, InProcessTransformResult result) {
-      Iterable<? extends CommittedBundle<?>> resultBundles =
+      CommittedResult committedResult =
           evaluationContext.handleResult(inputBundle, Collections.<TimerData>emptyList(), result);
-      for (CommittedBundle<?> outputBundle : resultBundles) {
+      for (CommittedBundle<?> outputBundle : committedResult.getOutputs()) {
         allUpdates.offer(ExecutorUpdate.fromBundle(outputBundle));
       }
-      return resultBundles;
+      return committedResult;
     }
 
     @Override
@@ -246,14 +246,14 @@ private TimerCompletionCallback(Iterable<TimerData> timers) {
     }
 
     @Override
-    public Iterable<? extends CommittedBundle<?>> handleResult(
+    public CommittedResult handleResult(
         CommittedBundle<?> inputBundle, InProcessTransformResult result) {
-      Iterable<? extends CommittedBundle<?>> resultBundles =
+      CommittedResult committedResult =
           evaluationContext.handleResult(inputBundle, timers, result);
-      for (CommittedBundle<?> outputBundle : resultBundles) {
+      for (CommittedBundle<?> outputBundle : committedResult.getOutputs()) {
         allUpdates.offer(ExecutorUpdate.fromBundle(outputBundle));
       }
-      return resultBundles;
+      return committedResult;
     }
 
     @Override

sdks/java/core/src/main/java/org/apache/beam/sdk/runners/inprocess/InProcessEvaluationContext.java

@@ -145,7 +145,7 @@ private InProcessEvaluationContext(
    * @param result the result of evaluating the input bundle
    * @return the committed bundles contained within the handled {@code result}
    */
-  public synchronized Iterable<? extends CommittedBundle<?>> handleResult(
+  public synchronized CommittedResult handleResult(
       @Nullable CommittedBundle<?> completedBundle,
       Iterable<TimerData> completedTimers,
       InProcessTransformResult result) {
@@ -176,7 +176,7 @@ public synchronized Iterable<? extends CommittedBundle<?>> handleResult(
         applicationStateInternals.remove(stepAndKey);
       }
     }
-    return committedBundles;
+    return CommittedResult.create(result, committedBundles);
   }
 
   private Iterable<? extends CommittedBundle<?>> commitBundles(

sdks/java/core/src/main/java/org/apache/beam/sdk/runners/inprocess/TransformExecutor.java

@@ -158,9 +158,9 @@ private InProcessTransformResult finishBundle(
       TransformEvaluator<T> evaluator, Collection<ModelEnforcement<T>> enforcements)
       throws Exception {
     InProcessTransformResult result = evaluator.finishBundle();
-    Iterable<? extends CommittedBundle<?>> outputs = onComplete.handleResult(inputBundle, result);
+    CommittedResult outputs = onComplete.handleResult(inputBundle, result);
     for (ModelEnforcement<T> enforcement : enforcements) {
-      enforcement.afterFinish(inputBundle, result, outputs);
+      enforcement.afterFinish(inputBundle, result, outputs.getOutputs());
     }
     return result;
   }

sdks/java/core/src/test/java/org/apache/beam/sdk/runners/inprocess/CommittedResultTest.java

@@ -0,0 +1,78 @@
+/*
+ * 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.sdk.runners.inprocess;
+
+import static org.junit.Assert.assertThat;
+
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.AppliedPTransform;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.util.WindowingStrategy;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+
+import com.google.common.collect.ImmutableList;
+
+import org.hamcrest.Matchers;
+import org.joda.time.Instant;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.io.Serializable;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Tests for {@link CommittedResult}.
+ */
+@RunWith(JUnit4.class)
+public class CommittedResultTest implements Serializable {
+  private transient TestPipeline p = TestPipeline.create();
+  private transient AppliedPTransform<?, ?, ?> transform =
+      AppliedPTransform.of("foo", p.begin(), PDone.in(p), new PTransform<PBegin, PDone>() {
+      });
+  private transient BundleFactory bundleFactory = InProcessBundleFactory.create();
+
+  @Test
+  public void getTransformExtractsFromResult() {
+    CommittedResult result =
+        CommittedResult.create(StepTransformResult.withoutHold(transform).build(),
+            Collections.<InProcessPipelineRunner.CommittedBundle<?>>emptyList());
+
+    assertThat(result.getTransform(), Matchers.<AppliedPTransform<?, ?, ?>>equalTo(transform));
+  }
+
+  @Test
+  public void getOutputsEqualInput() {
+    List<? extends InProcessPipelineRunner.CommittedBundle<?>> outputs =
+        ImmutableList.of(bundleFactory.createRootBundle(PCollection.createPrimitiveOutputInternal(p,
+            WindowingStrategy.globalDefault(),
+            PCollection.IsBounded.BOUNDED)).commit(Instant.now()),
+            bundleFactory.createRootBundle(PCollection.createPrimitiveOutputInternal(p,
+                WindowingStrategy.globalDefault(),
+                PCollection.IsBounded.UNBOUNDED)).commit(Instant.now()));
+    CommittedResult result =
+        CommittedResult.create(StepTransformResult.withoutHold(transform).build(), outputs);
+
+    assertThat(result.getOutputs(), Matchers.containsInAnyOrder(outputs.toArray()));
+  }
+}

sdks/java/core/src/test/java/org/apache/beam/sdk/runners/inprocess/InProcessEvaluationContextTest.java

@@ -460,7 +460,7 @@ public void isDoneWithPartiallyDone() {
 
     UncommittedBundle<Integer> rootBundle = context.createRootBundle(created);
     rootBundle.add(WindowedValue.valueInGlobalWindow(1));
-    Iterable<? extends CommittedBundle<?>> handleResult =
+    CommittedResult handleResult =
         context.handleResult(
             null,
             ImmutableList.<TimerData>of(),
@@ -469,7 +469,7 @@ public void isDoneWithPartiallyDone() {
                 .build());
     @SuppressWarnings("unchecked")
     CommittedBundle<Integer> committedBundle =
-        (CommittedBundle<Integer>) Iterables.getOnlyElement(handleResult);
+        (CommittedBundle<Integer>) Iterables.getOnlyElement(handleResult.getOutputs());
     context.handleResult(
         null,
         ImmutableList.<TimerData>of(),

sdks/java/core/src/test/java/org/apache/beam/sdk/runners/inprocess/TransformExecutorTest.java

@@ -487,11 +487,11 @@ private RegisteringCompletionCallback(CountDownLatch onMethod) {
     }
 
     @Override
-    public Iterable<? extends CommittedBundle<?>> handleResult(
+    public CommittedResult handleResult(
         CommittedBundle<?> inputBundle, InProcessTransformResult result) {
       handledResult = result;
       onMethod.countDown();
-      return Collections.emptyList();
+      return CommittedResult.create(result, Collections.<CommittedBundle<?>>emptyList());
     }
 
     @Override