Merge pull request #13523: Simplify LateDataDropping runner.

je-ik · web-flow · commit 3640ae23310b · 2021-02-12T09:13:51.000+01:00
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java
@@ -17,8 +17,8 @@
  */
 package org.apache.beam.runners.core;
 
-import java.util.stream.Collectors;
-import java.util.stream.StreamSupport;
+import java.util.ArrayList;
+import java.util.List;
 import org.apache.beam.sdk.metrics.Counter;
 import org.apache.beam.sdk.metrics.Metrics;
 import org.apache.beam.sdk.state.TimeDomain;
@@ -29,7 +29,6 @@
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
 import org.joda.time.Instant;
 
 /**
@@ -124,52 +123,37 @@ public LateDataFilter(
      */
     public <K, InputT> Iterable<WindowedValue<InputT>> filter(
         final K key, Iterable<WindowedValue<InputT>> elements) {
-      Iterable<Iterable<WindowedValue<InputT>>> windowsExpandedElements =
-          StreamSupport.stream(elements.spliterator(), false)
-              .map(
-                  input ->
-                      input.getWindows().stream()
-                          .map(
-                              window ->
-                                  WindowedValue.of(
-                                      input.getValue(),
-                                      input.getTimestamp(),
-                                      window,
-                                      input.getPane()))
-                          .collect(Collectors.toList()))
-              .collect(Collectors.toList());
-      Iterable<WindowedValue<InputT>> concatElements = Iterables.concat(windowsExpandedElements);
-
-      // Bump the counter separately since we don't want multiple iterations to
-      // increase it multiple times.
-      for (WindowedValue<InputT> input : concatElements) {
-        BoundedWindow window = Iterables.getOnlyElement(input.getWindows());
-        if (canDropDueToExpiredWindow(window)) {
-          // The element is too late for this window.
-          droppedDueToLateness.inc();
-          WindowTracing.debug(
-              "{}: Dropping element at {} for key:{}; window:{} "
-                  + "since too far behind inputWatermark:{}; outputWatermark:{}",
-              LateDataFilter.class.getSimpleName(),
-              input.getTimestamp(),
-              key,
-              window,
-              timerInternals.currentInputWatermarkTime(),
-              timerInternals.currentOutputWatermarkTime());
+      final List<WindowedValue<InputT>> nonLateElements = new ArrayList<>();
+      for (WindowedValue<InputT> element : elements) {
+        for (BoundedWindow window : element.getWindows()) {
+          if (canDropDueToExpiredWindow(window)) {
+            // The element is too late for this window.
+            droppedDueToLateness.inc();
+            WindowTracing.debug(
+                "{}: Dropping element at {} for key:{}; window:{} "
+                    + "since too far behind inputWatermark:{}; outputWatermark:{}",
+                LateDataFilter.class.getSimpleName(),
+                element.getTimestamp(),
+                key,
+                window,
+                timerInternals.currentInputWatermarkTime(),
+                timerInternals.currentOutputWatermarkTime());
+          } else {
+            nonLateElements.add(
+                WindowedValue.of(
+                    element.getValue(), element.getTimestamp(), window, element.getPane()));
+          }
         }
       }
-
-      // return nonLateElements
-      return StreamSupport.stream(concatElements.spliterator(), false)
-          .filter(
-              input -> {
-                BoundedWindow window = Iterables.getOnlyElement(input.getWindows());
-                return !canDropDueToExpiredWindow(window);
-              })
-          .collect(Collectors.toList());
+      return nonLateElements;
     }
 
-    /** Is {@code window} expired w.r.t. the garbage collection watermark? */
+    /**
+     * Is {@code window} expired w.r.t. the garbage collection watermark?
+     *
+     * @param window Window to check for expiration.
+     * @return True if element can be dropped.
+     */
     private boolean canDropDueToExpiredWindow(BoundedWindow window) {
       Instant inputWM = timerInternals.currentInputWatermarkTime();
       return LateDataUtils.garbageCollectionTime(window, windowingStrategy).isBefore(inputWM);
