fix npe in ReduceFnRunner, add additional tests, add timestamp validation.

Author: stankiewicz

runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java

@@ -472,23 +472,22 @@ public <T> void outputWindowedValue(
               element.causedByDrain()));
     }
 
-    @Override
-    public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {
+    private void noteOutputAndObserveTimestamp(Instant timestamp) {
       noteOutput();
       if (watermarkEstimator instanceof TimestampObservingWatermarkEstimator) {
-        ((TimestampObservingWatermarkEstimator) watermarkEstimator)
-            .observeTimestamp(windowedValue.getTimestamp());
+        ((TimestampObservingWatermarkEstimator) watermarkEstimator).observeTimestamp(timestamp);
       }
+    }
+
+    @Override
+    public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {
+      noteOutputAndObserveTimestamp(windowedValue.getTimestamp());
       outputReceiver.output(mainOutputTag, windowedValue);
     }
 
     @Override
     public <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedValue) {
-      noteOutput();
-      if (watermarkEstimator instanceof TimestampObservingWatermarkEstimator) {
-        ((TimestampObservingWatermarkEstimator) watermarkEstimator)
-            .observeTimestamp(windowedValue.getTimestamp());
-      }
+      noteOutputAndObserveTimestamp(windowedValue.getTimestamp());
       outputReceiver.output(tag, windowedValue);
     }
 

runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java

@@ -376,7 +376,7 @@ public void processElements(Iterable<WindowedValue<InputT>> values) throws Excep
       emit(
           contextFactory.base(window, StateStyle.DIRECT),
           contextFactory.base(window, StateStyle.RENAMED),
-          null);
+          CausedByDrain.NORMAL);
     }
 
     // We're all done with merging and emitting elements so can compress the activeWindow state.

runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java

@@ -445,11 +445,13 @@ public <T> void output(TupleTag<T> tag, T output) {
 
     @Override
     public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {
+      checkTimestamp(elem.getTimestamp(), windowedValue.getTimestamp());
       SimpleDoFnRunner.this.outputWindowedValue(mainOutputTag, windowedValue);
     }
 
     @Override
     public <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedValue) {
+      checkTimestamp(elem.getTimestamp(), windowedValue.getTimestamp());
       SimpleDoFnRunner.this.outputWindowedValue(tag, windowedValue);
     }
 
@@ -1039,11 +1041,13 @@ public <T> void outputWindowedValue(
 
     @Override
     public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {
+      checkTimestamp(timestamp(), windowedValue.getTimestamp());
       SimpleDoFnRunner.this.outputWindowedValue(mainOutputTag, windowedValue);
     }
 
     @Override
     public <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedValue) {
+      checkTimestamp(timestamp(), windowedValue.getTimestamp());
       SimpleDoFnRunner.this.outputWindowedValue(tag, windowedValue);
     }
 
@@ -1308,11 +1312,13 @@ public <T> void outputWindowedValue(
 
     @Override
     public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {
+      checkTimestamp(this.timestamp, windowedValue.getTimestamp());
       SimpleDoFnRunner.this.outputWindowedValue(mainOutputTag, windowedValue);
     }
 
     @Override
     public <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedValue) {
+      checkTimestamp(this.timestamp, windowedValue.getTimestamp());
       SimpleDoFnRunner.this.outputWindowedValue(tag, windowedValue);
     }
 

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java

@@ -644,10 +644,9 @@ public <T> void outputWithTimestamp(TupleTag<T> tag, T output, Instant timestamp
                   CausedByDrain.NORMAL));
     }
 
-    @Override
-    public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {
+    private <T> void doOutputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedValue) {
       for (BoundedWindow w : windowedValue.getWindows()) {
-        getMutableOutput(mainOutputTag)
+        getMutableOutput(tag)
             .add(
                 ValueInSingleWindow.of(
                     windowedValue.getValue(),
@@ -660,20 +659,14 @@ public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {
       }
     }
 
+    @Override
+    public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {
+      doOutputWindowedValue(mainOutputTag, windowedValue);
+    }
+
     @Override
     public <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedValue) {
-      for (BoundedWindow w : windowedValue.getWindows()) {
-        getMutableOutput(tag)
-            .add(
-                ValueInSingleWindow.of(
-                    windowedValue.getValue(),
-                    windowedValue.getTimestamp(),
-                    w,
-                    windowedValue.getPaneInfo(),
-                    windowedValue.getRecordId(),
-                    windowedValue.getRecordOffset(),
-                    windowedValue.causedByDrain()));
-      }
+      doOutputWindowedValue(tag, windowedValue);
     }
 
     @Override

sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java

@@ -322,6 +322,7 @@ public static <T> WindowedValue<T> of(
       @Nullable Long currentRecordOffset,
       CausedByDrain causedByDrain) {
     checkArgument(paneInfo != null, "WindowedValue requires PaneInfo, but it was null");
+    checkArgument(causedByDrain != null, "WindowedValue requires causedByDrain, but it was null");
 
     boolean isGlobal = GlobalWindow.INSTANCE.equals(window);
     if (isGlobal && BoundedWindow.TIMESTAMP_MIN_VALUE.equals(timestamp)) {
@@ -861,6 +862,10 @@ public static void setMetadataSupported() {
       metadataSupported = true;
     }
 
+    public static void setMetadataNotSupported() {
+      metadataSupported = false;
+    }
+
     public static boolean isMetadataSupported() {
       return metadataSupported;
     }

sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/MetadataPropagationTest.java

@@ -23,6 +23,8 @@
 import org.apache.beam.sdk.values.CausedByDrain;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.WindowedValues;
+import org.junit.After;
+import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
@@ -70,6 +72,27 @@ public void testMetadataPropagationAcrossShuffleParameter() {
       pipeline.run();
     }
 
+    @After
+    public void turnOffMetadata() {
+      WindowedValues.WindowedValueCoder.setMetadataNotSupported();
+    }
+
+    @Test
+    @Category(NeedsRunner.class)
+    public void testDefaultMetadataPropagationAcrossShuffleParameter() {
+      Assert.assertFalse(WindowedValues.WindowedValueCoder.isMetadataSupported());
+      PCollection<String> results =
+          pipeline
+              .apply(Create.of(1))
+              .apply(ParDo.of(new CausedByDrainSettingDoFn()))
+              .apply(Redistribute.arbitrarily())
+              .apply(ParDo.of(new CausedByDrainExtractingDoFn()));
+
+      PAssert.that(results).containsInAnyOrder("NORMAL");
+
+      pipeline.run();
+    }
+
     @Test
     @Category(NeedsRunner.class)
     public void testMetadataPropagationParameter() {

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java

@@ -1954,11 +1954,13 @@ public <T> void outputWindowedValue(
 
     @Override
     public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {
+      checkTimestamp(windowedValue.getTimestamp());
       outputTo(mainOutputConsumer, windowedValue);
     }
 
     @Override
     public <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedValue) {
+      checkTimestamp(windowedValue.getTimestamp());
       outputTo((FnDataReceiver) localNameToConsumer.get(tag.getId()), windowedValue);
     }