spotless

Author: stankiewicz

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

@@ -478,8 +478,6 @@ public String getErrorContext() {
         watermarkEstimatorStateT = watermarkEstimatorState.read();
       }
 
-      // elementAndRestriction should be patched with drain info
-
       final WatermarkEstimator<WatermarkEstimatorStateT> watermarkEstimator =
           invoker.invokeNewWatermarkEstimator(
               new BaseArgumentProvider<InputT, OutputT>() {
@@ -613,14 +611,14 @@ public String getErrorContext() {
           timerInternals.currentProcessingTime().plus(result.getContinuation().resumeDelay());
       holdState.add(futureOutputWatermark);
       // Set a timer to continue processing this element, but only when no drain
-      if(timer.causedByDrain() == CausedByDrain.NORMAL) {
+      if (timer == null || timer.causedByDrain() == CausedByDrain.NORMAL) {
         timerInternals.setTimer(
-          TimerInternals.TimerData.of(
-            stateNamespace,
-            wakeupTime,
-            wakeupTime,
-            TimeDomain.PROCESSING_TIME,
-            timer == null ? CausedByDrain.NORMAL : timer.causedByDrain()));
+            TimerInternals.TimerData.of(
+                stateNamespace,
+                wakeupTime,
+                wakeupTime,
+                TimeDomain.PROCESSING_TIME,
+                CausedByDrain.NORMAL));
       }
     }
 

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

@@ -1804,11 +1804,6 @@ public <T> void outputWindowedValue(
       outputTo(consumer, WindowedValues.of(output, timestamp, windows, paneInfo));
     }
 
-    @Override
-    public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
-      return currentElement.causedByDrain();
-    }
-
     @Override
     public State state(String stateId, boolean alwaysFetched) {
       StateDeclaration stateDeclaration = doFnSignature.stateDeclarations().get(stateId);
@@ -1860,11 +1855,6 @@ public TimerMap timerFamily(String timerFamilyId) {
           currentElement.getTimestamp(),
           currentElement.getPaneInfo());
     }
-
-    @Override
-    public CausedByDrain causedByDrain() {
-      return currentElement.causedByDrain();
-    }
   }
 
   /** Provides arguments for a {@link DoFnInvoker} for a non-window observing method. */
@@ -1946,16 +1936,6 @@ public <T> void outputWindowedValue(
       }
       outputTo(consumer, WindowedValues.of(output, timestamp, windows, paneInfo));
     }
-
-    @Override
-    public CausedByDrain causedByDrain() {
-      return currentElement.causedByDrain();
-    }
-
-    @Override
-    public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
-      return currentElement.causedByDrain();
-    }
   }
 
   /** Provides base arguments for a {@link DoFnInvoker} for a non-window observing method. */
@@ -2254,6 +2234,16 @@ public Object watermarkEstimatorState() {
     public WatermarkEstimator<?> watermarkEstimator() {
       return currentWatermarkEstimator;
     }
+
+    @Override
+    public CausedByDrain causedByDrain() {
+      return currentElement.causedByDrain();
+    }
+
+    @Override
+    public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
+      return currentElement.causedByDrain();
+    }
   }
 
   /**