Continues unifying ReaderIterator and Source.Reader

* Flattens the ReaderIterator hierarchy a bit:
  removes Abstract{Bounded,}ReaderIterator, instead merges them all
  into LegacyReaderIterator directly.
* Converts InMemoryReaderIterator, PubsubReaderIterator,
  UngroupedWindmillReaderIterator, ConcatIterator, and a couple of
  anonymous classes to the new interface.
  File-based and shuffle readers remain.
* Extracts a common base class from Pubsub and UngroupedWindmill readers.
* Converts most call sites that could use Source.Reader instead of
  ReaderIterator to use either Source.Reader, or the concrete class
  of the reader, to ease future conversion.

----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=112728940

Author: jkff

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/AvroByteReader.java

@@ -21,7 +21,6 @@
 import com.google.cloud.dataflow.sdk.options.PipelineOptions;
 import com.google.cloud.dataflow.sdk.util.CoderUtils;
 import com.google.cloud.dataflow.sdk.util.WindowedValue;
-import com.google.cloud.dataflow.sdk.util.common.worker.AbstractBoundedReaderIterator;
 import com.google.cloud.dataflow.sdk.util.common.worker.NativeReader;
 
 import org.apache.avro.Schema;
@@ -51,11 +50,11 @@ public AvroByteReader(String filename, @Nullable Long startPosition, @Nullable L
   }
 
   @Override
-  public NativeReaderIterator<T> iterator() throws IOException {
+  public AvroByteFileIterator iterator() throws IOException {
     return new AvroByteFileIterator();
   }
 
-  class AvroByteFileIterator extends AbstractBoundedReaderIterator<T> {
+  class AvroByteFileIterator extends LegacyReaderIterator<T> {
     private final LegacyReaderIterator<WindowedValue<ByteBuffer>> avroFileIterator;
 
     public AvroByteFileIterator() throws IOException {

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/AvroReader.java

@@ -26,7 +26,6 @@
 import com.google.cloud.dataflow.sdk.io.OffsetBasedSource;
 import com.google.cloud.dataflow.sdk.options.PipelineOptions;
 import com.google.cloud.dataflow.sdk.util.WindowedValue;
-import com.google.cloud.dataflow.sdk.util.common.worker.AbstractBoundedReaderIterator;
 import com.google.cloud.dataflow.sdk.util.common.worker.NativeReader;
 
 import org.apache.avro.generic.GenericRecord;
@@ -78,7 +77,7 @@ public AvroReader(String filename, @Nullable Long startPosition, @Nullable Long
   }
 
   @Override
-  public LegacyReaderIterator<WindowedValue<T>> iterator() throws IOException {
+  public AvroFileIterator iterator() throws IOException {
     Long endPosition = this.endPosition;
     Long startPosition = this.startPosition;
     if (endPosition == null) {
@@ -99,7 +98,7 @@ public LegacyReaderIterator<WindowedValue<T>> iterator() throws IOException {
     return new AvroFileIterator((AvroSource.AvroReader<T>) reader);
   }
 
-  class AvroFileIterator extends AbstractBoundedReaderIterator<WindowedValue<T>> {
+  class AvroFileIterator extends LegacyReaderIterator<WindowedValue<T>> {
     final AvroSource.AvroReader<T> reader;
     boolean hasStarted = false;
     long blockOffset = -1;

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/BigQueryReader.java

@@ -25,7 +25,6 @@
 import com.google.cloud.dataflow.sdk.util.BigQueryTableRowIterator;
 import com.google.cloud.dataflow.sdk.util.Transport;
 import com.google.cloud.dataflow.sdk.util.WindowedValue;
-import com.google.cloud.dataflow.sdk.util.common.worker.AbstractBoundedReaderIterator;
 import com.google.cloud.dataflow.sdk.util.common.worker.NativeReader;
 import com.google.common.annotations.VisibleForTesting;
 
@@ -113,8 +112,7 @@ public BigQueryReaderIterator iterator() throws IOException {
    * A ReaderIterator that yields TableRow objects for each row of a BigQuery table.
    */
   @VisibleForTesting
-  static class BigQueryReaderIterator
-      extends AbstractBoundedReaderIterator<WindowedValue<TableRow>> {
+  static class BigQueryReaderIterator extends LegacyReaderIterator<WindowedValue<TableRow>> {
     private BigQueryTableRowIterator rowIterator;
 
     public BigQueryReaderIterator(TableReference tableRef, Bigquery bigQueryClient) {

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/ConcatReader.java

@@ -30,8 +30,8 @@
 import com.google.cloud.dataflow.sdk.runners.worker.SourceTranslationUtils.DataflowReaderProgress;
 import com.google.cloud.dataflow.sdk.util.ExecutionContext;
 import com.google.cloud.dataflow.sdk.util.common.CounterSet;
-import com.google.cloud.dataflow.sdk.util.common.worker.AbstractBoundedReaderIterator;
 import com.google.cloud.dataflow.sdk.util.common.worker.NativeReader;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 
 import org.slf4j.Logger;
@@ -42,6 +42,8 @@
 import java.util.List;
 import java.util.NoSuchElementException;
 
+import javax.annotation.Nullable;
+
 /**
  * A {@link NativeReader} that reads elements from a given set of encoded {@link Source}s. Creates
  * {@link NativeReader}s for sources lazily, i.e. only when elements from the particular
@@ -96,7 +98,7 @@ public Iterator<Source> getSources() {
   }
 
   @Override
-  public LegacyReaderIterator<T> iterator() throws IOException {
+  public ConcatIterator<T> iterator() throws IOException {
     return new ConcatIterator<T>(
         registry,
         options,
@@ -106,16 +108,16 @@ public LegacyReaderIterator<T> iterator() throws IOException {
         sources);
   }
 
-  private static class ConcatIterator<T> extends AbstractBoundedReaderIterator<T> {
+  @VisibleForTesting
+  static class ConcatIterator<T> extends NativeReaderIterator<T> {
     private int currentIteratorIndex = -1;
-    private LegacyReaderIterator<T> currentIterator = null;
+    @Nullable private NativeReaderIterator<T> currentIterator = null;
     private final List<Source> sources;
     private final PipelineOptions options;
     private final ExecutionContext executionContext;
     private final CounterSet.AddCounterMutator addCounterMutator;
     private final String operationName;
     private final OffsetRangeTracker rangeTracker;
-    private boolean isAtFirstRecordInCurrentSource = true;
     private final ReaderFactory.Registry registry;
 
     public ConcatIterator(
@@ -135,16 +137,29 @@ public ConcatIterator(
     }
 
     @Override
-    protected boolean hasNextImpl() throws IOException {
-      for (;;) {
-        if (currentIterator != null && currentIterator.hasNext()) {
-          break;
-        }
+    public boolean start() throws IOException {
+      return advance();
+    }
 
+    @Override
+    public boolean advance() throws IOException {
+      while (true) {
+        // Invariant: we call currentIterator.start() immediately when opening an iterator
+        // (below). So if currentIterator != null, then start() has already been called on it.
+        if (currentIterator != null && currentIterator.advance()) {
+          // Happy case: current iterator has a next record.
+          return true;
+        }
+        // Now current iterator is either non-existent or exhausted.
+        // Close it, and try opening a new one.
         if (currentIterator != null) {
           currentIterator.close();
+          currentIterator = null;
         }
 
+        if (!rangeTracker.tryReturnRecordAt(true, currentIteratorIndex + 1)) {
+          return false;
+        }
         currentIteratorIndex++;
         if (currentIteratorIndex == sources.size()) {
           // All sources were read.
@@ -158,20 +173,26 @@ protected boolean hasNextImpl() throws IOException {
               (NativeReader<T>)
                   registry.create(
                       currentSource, options, executionContext, addCounterMutator, operationName);
-          currentIterator = (LegacyReaderIterator) currentReader.iterator();
-          isAtFirstRecordInCurrentSource = true;
+          currentIterator = currentReader.iterator();
         } catch (Exception e) {
           throw new IOException("Failed to create a reader for source: " + currentSource, e);
         }
+        if (!currentIterator.start()) {
+          currentIterator.close();
+          currentIterator = null;
+          continue;
+        }
+        // Happy case: newly opened iterator has a first record.
+        return true;
       }
-
-      return rangeTracker.tryReturnRecordAt(isAtFirstRecordInCurrentSource, currentIteratorIndex);
     }
 
     @Override
-    protected T nextImpl() throws IOException, NoSuchElementException {
-      isAtFirstRecordInCurrentSource = false;
-      return currentIterator.next();
+    public T getCurrent() throws NoSuchElementException {
+      if (currentIterator == null) {
+        throw new NoSuchElementException();
+      }
+      return currentIterator.getCurrent();
     }
 
     @Override

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/GroupingShuffleReader.java

@@ -40,7 +40,6 @@
 import com.google.cloud.dataflow.sdk.util.common.ElementByteSizeObservableIterator;
 import com.google.cloud.dataflow.sdk.util.common.Reiterable;
 import com.google.cloud.dataflow.sdk.util.common.Reiterator;
-import com.google.cloud.dataflow.sdk.util.common.worker.AbstractBoundedReaderIterator;
 import com.google.cloud.dataflow.sdk.util.common.worker.BatchingShuffleEntryReader;
 import com.google.cloud.dataflow.sdk.util.common.worker.GroupingShuffleEntryIterator;
 import com.google.cloud.dataflow.sdk.util.common.worker.KeyGroupedShuffleEntries;
@@ -193,7 +192,7 @@ final GroupingShuffleReaderIterator<K, V> iterator(ShuffleEntryReader reader) {
    */
   @VisibleForTesting
   static final class GroupingShuffleReaderIterator<K, V>
-      extends AbstractBoundedReaderIterator<WindowedValue<KV<K, Reiterable<V>>>> {
+      extends LegacyReaderIterator<WindowedValue<KV<K, Reiterable<V>>>> {
     // The enclosing GroupingShuffleReader.
     private final GroupingShuffleReader<K, V> parentReader;
 

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/InMemoryReader.java

@@ -20,22 +20,25 @@
 import static com.google.cloud.dataflow.sdk.runners.worker.SourceTranslationUtils.cloudPositionToReaderPosition;
 import static com.google.cloud.dataflow.sdk.runners.worker.SourceTranslationUtils.cloudProgressToReaderProgress;
 import static com.google.cloud.dataflow.sdk.runners.worker.SourceTranslationUtils.splitRequestToApproximateSplitRequest;
-import static java.lang.Math.min;
+import static com.google.common.base.MoreObjects.firstNonNull;
+import static com.google.common.base.Preconditions.checkArgument;
 
 import com.google.api.services.dataflow.model.ApproximateReportedProgress;
 import com.google.cloud.dataflow.sdk.coders.Coder;
+import com.google.cloud.dataflow.sdk.coders.CoderException;
 import com.google.cloud.dataflow.sdk.io.range.OffsetRangeTracker;
 import com.google.cloud.dataflow.sdk.util.CoderUtils;
 import com.google.cloud.dataflow.sdk.util.StringUtils;
-import com.google.cloud.dataflow.sdk.util.common.worker.AbstractBoundedReaderIterator;
 import com.google.cloud.dataflow.sdk.util.common.worker.NativeReader;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 import java.util.List;
+import java.util.NoSuchElementException;
 
 import javax.annotation.Nullable;
 
@@ -52,26 +55,20 @@ public class InMemoryReader<T> extends NativeReader<T> {
   final int endIndex;
   final Coder<T> coder;
 
-  public InMemoryReader(List<String> encodedElements, @Nullable Long startIndex,
-      @Nullable Long endIndex, Coder<T> coder) {
+  public InMemoryReader(
+      List<String> encodedElements,
+      @Nullable Integer startIndex,
+      @Nullable Integer endIndex,
+      Coder<T> coder) {
+    checkNotNull(encodedElements);
     this.encodedElements = encodedElements;
     int maxIndex = encodedElements.size();
-    if (startIndex == null) {
-      this.startIndex = 0;
-    } else {
-      if (startIndex < 0) {
-        throw new IllegalArgumentException("start index should be >= 0");
-      }
-      this.startIndex = (int) min(startIndex, maxIndex);
-    }
-    if (endIndex == null) {
-      this.endIndex = maxIndex;
-    } else {
-      if (endIndex < this.startIndex) {
-        throw new IllegalArgumentException("end index should be >= start index");
-      }
-      this.endIndex = (int) min(endIndex, maxIndex);
-    }
+    this.startIndex = Math.min(maxIndex, firstNonNull(startIndex, 0));
+    this.endIndex = Math.min(maxIndex, firstNonNull(endIndex, maxIndex));
+    checkArgument(this.startIndex >= 0, "negative start index: " + startIndex);
+    checkArgument(
+        this.endIndex >= this.startIndex,
+        "end index before start: [" + this.startIndex + ", " + this.endIndex + ")");
     this.coder = coder;
   }
 
@@ -88,24 +85,45 @@ public double getTotalParallelism() {
   /**
    * A ReaderIterator that yields an in-memory list of elements.
    */
-  class InMemoryReaderIterator extends AbstractBoundedReaderIterator<T> {
+  class InMemoryReaderIterator extends NativeReaderIterator<T> {
     @VisibleForTesting
     OffsetRangeTracker tracker;
-    private int nextIndex;
+    @Nullable private Integer lastReturnedIndex;
+    private T current;
 
     public InMemoryReaderIterator() {
       this.tracker = new OffsetRangeTracker(startIndex, endIndex);
-      this.nextIndex = startIndex;
+      this.lastReturnedIndex = null;
     }
 
     @Override
-    protected boolean hasNextImpl() {
-      return tracker.tryReturnRecordAt(true, nextIndex);
+    public boolean start() throws IOException {
+      Preconditions.checkState(lastReturnedIndex == null, "Already started");
+      if (!tracker.tryReturnRecordAt(true, startIndex)) {
+        return false;
+      }
+      current = decode(encodedElements.get(startIndex));
+      lastReturnedIndex = startIndex;
+      return true;
     }
 
     @Override
-    protected T nextImpl() throws IOException {
-      String encodedElementString = encodedElements.get(nextIndex++);
+    public boolean advance() throws IOException {
+      Preconditions.checkNotNull(lastReturnedIndex, "Not started");
+      if (!tracker.tryReturnRecordAt(true, (long) lastReturnedIndex + 1)) {
+        return false;
+      }
+      ++lastReturnedIndex;
+      current = decode(encodedElements.get(lastReturnedIndex));
+      return true;
+    }
+
+    @Override
+    public T getCurrent() throws NoSuchElementException {
+      return current;
+    }
+
+    private T decode(String encodedElementString) throws CoderException {
       // TODO: Replace with the real encoding used by the
       // front end, when we know what it is.
       byte[] encodedElement = StringUtils.jsonStringToByteArray(encodedElementString);
@@ -115,12 +133,15 @@ protected T nextImpl() throws IOException {
 
     @Override
     public Progress getProgress() {
+      if (lastReturnedIndex == null) {
+        return null;
+      }
       // Currently we assume that only a record index position is reported as
       // current progress. An implementer can override this method to update
       // other metrics, e.g. completion percentage or remaining time.
       com.google.api.services.dataflow.model.Position currentPosition =
           new com.google.api.services.dataflow.model.Position();
-      currentPosition.setRecordIndex((long) nextIndex);
+      currentPosition.setRecordIndex((long) lastReturnedIndex);
 
       ApproximateReportedProgress progress = new ApproximateReportedProgress();
       progress.setPosition(currentPosition);
@@ -130,7 +151,8 @@ public Progress getProgress() {
 
     @Override
     public double getRemainingParallelism() {
-      return tracker.getStopPosition() - nextIndex;
+      // Use the starting index if no elements have yet been returned.
+      return tracker.getStopPosition() - firstNonNull(lastReturnedIndex, startIndex);
     }
 
     @Override

sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/worker/InMemoryReaderFactory.java

@@ -16,7 +16,7 @@
 
 package com.google.cloud.dataflow.sdk.runners.worker;
 
-import static com.google.cloud.dataflow.sdk.util.Structs.getLong;
+import static com.google.cloud.dataflow.sdk.util.Structs.getInt;
 import static com.google.cloud.dataflow.sdk.util.Structs.getStrings;
 
 import com.google.cloud.dataflow.sdk.coders.Coder;
@@ -51,7 +51,8 @@ public NativeReader<?> create(
   <T> InMemoryReader<T> create(CloudObject spec, Coder<T> coder) throws Exception {
     return new InMemoryReader<>(
         getStrings(spec, PropertyNames.ELEMENTS, Collections.<String>emptyList()),
-        getLong(spec, PropertyNames.START_INDEX, null),
-        getLong(spec, PropertyNames.END_INDEX, null), coder);
+        getInt(spec, PropertyNames.START_INDEX, null),
+        getInt(spec, PropertyNames.END_INDEX, null),
+        coder);
   }
 }