Rolling back due to an internal test failure.

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

Author: chamikaramj

google/cloud/dataflow/worker/executor_test.py

@@ -86,7 +86,7 @@ def __init__(self, source):
   def get_progress(self):
     next_progress = super(ProgressRequestRecordingInMemoryReader,
                           self).get_progress()
-    self.progress_record.append(next_progress.position.record_index)
+    self.progress_record.append(next_progress.percent_complete)
     return next_progress
 
 
@@ -351,7 +351,11 @@ def test_in_memory_source_progress_reporting(self):
     ]))
     self.assertEqual(elements, output_buffer)
 
-    expected_progress_record = range(len(elements))
+    expected_progress_record = []
+    len_elements = len(elements)
+    for i in range(len_elements):
+      expected_progress_record.append(float(i + 1) / len_elements)
+
     self.assertEqual(expected_progress_record,
                      source.last_reader.progress_record)
 

google/cloud/dataflow/worker/inmemory.py

@@ -15,11 +15,9 @@
 """In-memory input source."""
 
 import itertools
-import logging
 
 from google.cloud.dataflow import coders
 from google.cloud.dataflow.io import iobase
-from google.cloud.dataflow.io import range_trackers
 
 
 class InMemorySource(iobase.NativeSource):
@@ -55,14 +53,11 @@ class InMemoryReader(iobase.NativeSourceReader):
   """A reader for in-memory source."""
 
   def __init__(self, source):
-    self._source = source
+    self.source = source
 
-    # Index of the last item returned by InMemoryReader.
-    # Initialized to None.
-    self._current_index = None
-
-    self._range_tracker = range_trackers.OffsetRangeTracker(
-        self._source.start_index, self._source.end_index)
+    # Index of the next item to be read by the InMemoryReader.
+    # Starts at source.start_index.
+    self.current_index = source.start_index
 
   def __enter__(self):
     return self
@@ -71,47 +66,21 @@ def __exit__(self, exception_type, exception_value, traceback):
     pass
 
   def __iter__(self):
-    for value in itertools.islice(self._source.elements,
-                                  self._source.start_index,
-                                  self._source.end_index):
-      claimed = False
-      if self._current_index is None:
-        claimed = self._range_tracker.try_return_record_at(
-            True, self._source.start_index)
-      else:
-        claimed = self._range_tracker.try_return_record_at(
-            True, self._current_index + 1)
-
-      if claimed:
-        if self._current_index is None:
-          self._current_index = self._source.start_index
-        else:
-          self._current_index += 1
-
-        yield self._source.coder.decode(value)
-      else:
-        return
+    for value in itertools.islice(self.source.elements,
+                                  self.source.start_index,
+                                  self.source.end_index):
+      self.current_index += 1
+      yield self.source.coder.decode(value)
 
   def get_progress(self):
-    if self._current_index is None:
-      return None
-
-    return iobase.ReaderProgress(
-        position=iobase.ReaderPosition(record_index=self._current_index))
-
-  def request_dynamic_split(self, dynamic_split_request):
-    assert dynamic_split_request is not None
-    progress = dynamic_split_request.progress
-    split_position = progress.position
-    if split_position is None:
-      logging.debug('InMemory reader only supports split requests that are '
-                    'based on positions. Received : %r', dynamic_split_request)
-
-    index_position = split_position.record_index
-    if index_position is None:
-      logging.debug('InMemory reader only supports split requests that are '
-                    'based on index positions. Received : %r',
-                    dynamic_split_request)
-
-    if self._range_tracker.try_split_at_position(index_position):
-      return iobase.DynamicSplitResultWithPosition(split_position)
+    if (self.current_index >= self.source.end_index or
+        self.source.start_index >= self.source.end_index):
+      percent_complete = 1
+    elif self.current_index == self.source.start_index:
+      percent_complete = 0
+    else:
+      percent_complete = (
+          float(self.current_index - self.source.start_index) / (
+              self.source.end_index - self.source.start_index))
+
+    return iobase.ReaderProgress(percent_complete=percent_complete)

google/cloud/dataflow/worker/inmemory_test.py

@@ -17,7 +17,6 @@
 import logging
 import unittest
 
-from google.cloud.dataflow.io import iobase
 from google.cloud.dataflow.worker import inmemory
 
 
@@ -42,115 +41,31 @@ def test_norange(self):
   def test_in_memory_source_updates_progress_none(self):
     source = inmemory.InMemorySource([], coder=FakeCoder())
     with source.reader() as reader:
-      self.assertEqual(None, reader.get_progress())
+      self.assertEqual(1, reader.get_progress().percent_complete)
 
   def test_in_memory_source_updates_progress_one(self):
     source = inmemory.InMemorySource([1], coder=FakeCoder())
     with source.reader() as reader:
-      self.assertEqual(None, reader.get_progress())
+      self.assertEqual(0, reader.get_progress().percent_complete)
       i = 0
       for item in reader:
-        self.assertEqual(i, reader.get_progress().position.record_index)
-        self.assertEqual(11, item)
         i += 1
+        self.assertEqual(11, item)
+        self.assertEqual(1, reader.get_progress().percent_complete)
       self.assertEqual(1, i)
-      self.assertEqual(0, reader.get_progress().position.record_index)
+      self.assertEqual(1, reader.get_progress().percent_complete)
 
   def test_in_memory_source_updates_progress_many(self):
     source = inmemory.InMemorySource([1, 2, 3, 4, 5], coder=FakeCoder())
     with source.reader() as reader:
-      self.assertEqual(None, reader.get_progress())
+      self.assertEqual(0, reader.get_progress().percent_complete)
       i = 0
       for item in reader:
-        self.assertEqual(i, reader.get_progress().position.record_index)
-        self.assertEqual(11 + i, item)
         i += 1
+        self.assertEqual(i + 10, item)
+        self.assertEqual(float(i) / 5, reader.get_progress().percent_complete)
       self.assertEqual(5, i)
-      self.assertEqual(4, reader.get_progress().position.record_index)
-
-  def try_splitting_reader_at(self, reader, split_request, expected_response):
-    actual_response = reader.request_dynamic_split(split_request)
-
-    if expected_response is None:
-      self.assertIsNone(actual_response)
-    else:
-      self.assertIsNotNone(actual_response.stop_position)
-      self.assertIsInstance(actual_response.stop_position,
-                            iobase.ReaderPosition)
-      self.assertIsNotNone(actual_response.stop_position.record_index)
-      self.assertEqual(expected_response.stop_position.record_index,
-                       actual_response.stop_position.record_index)
-
-  def test_in_memory_source_dynamic_split(self):
-    source = inmemory.InMemorySource([10, 20, 30, 40, 50, 60],
-                                     coder=FakeCoder())
-
-    # Unstarted reader
-    with source.reader() as reader:
-      self.try_splitting_reader_at(
-          reader,
-          iobase.DynamicSplitRequest(
-              iobase.ReaderProgress(
-                  position=iobase.ReaderPosition(record_index=2))),
-          None)
-
-    # Proposed split position out of range
-    with source.reader() as reader:
-      reader_iter = iter(reader)
-      next(reader_iter)
-      self.try_splitting_reader_at(
-          reader,
-          iobase.DynamicSplitRequest(
-              iobase.ReaderProgress(
-                  position=iobase.ReaderPosition(record_index=-1))),
-          None)
-      self.try_splitting_reader_at(
-          reader,
-          iobase.DynamicSplitRequest(
-              iobase.ReaderProgress(
-                  position=iobase.ReaderPosition(record_index=10))),
-          None)
-
-    # Already read past proposed split position
-    with source.reader() as reader:
-      reader_iter = iter(reader)
-      next(reader_iter)
-      next(reader_iter)
-      next(reader_iter)
-      self.try_splitting_reader_at(
-          reader,
-          iobase.DynamicSplitRequest(
-              iobase.ReaderProgress(
-                  position=iobase.ReaderPosition(record_index=1))),
-          None)
-
-      self.try_splitting_reader_at(
-          reader,
-          iobase.DynamicSplitRequest(
-              iobase.ReaderProgress(
-                  position=iobase.ReaderPosition(record_index=2))),
-          None)
-
-    # Successful split
-    with source.reader() as reader:
-      reader_iter = iter(reader)
-      next(reader_iter)
-      self.try_splitting_reader_at(
-          reader,
-          iobase.DynamicSplitRequest(
-              iobase.ReaderProgress(
-                  position=iobase.ReaderPosition(record_index=4))),
-          iobase.DynamicSplitResultWithPosition(
-              stop_position=iobase.ReaderPosition(record_index=4)))
-
-      self.try_splitting_reader_at(
-          reader,
-          iobase.DynamicSplitRequest(
-              iobase.ReaderProgress(
-                  position=iobase.ReaderPosition(record_index=2))),
-          iobase.DynamicSplitResultWithPosition(
-              stop_position=iobase.ReaderPosition(record_index=2)))
-
+      self.assertEqual(1, reader.get_progress().percent_complete)
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)