Adds the base API for creating new sources.

Sources API consists of the classes 'BoundedSource' and 'RangeTracker'. Updated 'iobase.RangeTracker' and related classes to the new API.

Adds a good amount of documentation along with the API classes.

Adds support for reading sources using DirectRunner.

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

Author: chamikaramj

google/cloud/dataflow/io/fileio.py

@@ -586,8 +586,7 @@ def __exit__(self, exception_type, exception_value, traceback):
 
   def __iter__(self):
     while True:
-      if not self.range_tracker.try_return_record_at(
-          is_at_split_point=True,
+      if not self.range_tracker.try_claim(
           record_start=self.current_offset):
         # Reader has completed reading the set of records in its range. Note
         # that the end offset of the range may be smaller than the original
@@ -620,16 +619,15 @@ def request_dynamic_split(self, dynamic_split_request):
           return
         split_position = iobase.ReaderPosition()
         split_position.byte_offset = (
-            self.range_tracker.get_position_for_fraction_consumed(
-                percent_complete))
+            self.range_tracker.position_at_fraction(percent_complete))
       else:
         logging.warning(
             'TextReader requires either a position or a percentage of work to '
             'be complete to perform a dynamic split request. Requested: %r',
             dynamic_split_request)
         return
 
-    if self.range_tracker.try_split_at_position(split_position.byte_offset):
+    if self.range_tracker.try_split(split_position.byte_offset):
       return iobase.DynamicSplitResultWithPosition(split_position)
     else:
       return

google/cloud/dataflow/io/iobase.py

@@ -26,10 +26,13 @@
 the sink.
 """
 
+from collections import namedtuple
+
 import logging
 import uuid
 
 from google.cloud.dataflow import pvalue
+from google.cloud.dataflow.coders import PickleCoder
 from google.cloud.dataflow.pvalue import AsIter
 from google.cloud.dataflow.pvalue import AsSingleton
 from google.cloud.dataflow.transforms import core
@@ -295,17 +298,125 @@ def Write(self, o):  # pylint: disable=invalid-name
     raise NotImplementedError
 
 
-class RangeTracker(object):
-  """A thread-safe helper object for implementing dynamic work rebalancing.
+# Encapsulates information about a bundle of a source generated when method
+# BoundedSource.split() is invoked.
+# This is a named 4-tuple that has following fields.
+# * weight - a number that represents the size of the bundle. This value will
+#            be used to compare the relative sizes of bundles generated by the
+#            current source.
+#            The weight returned here could be specified using a unit of your
+#            choice (for example, bundles of sizes 100MB, 200MB, and 700MB may
+#            specify weights 100, 200, 700 or 1, 2, 7) but all bundles of a
+#            source should specify the weight using the same unit.
+# * source - a BoundedSource object for the  bundle.
+# * start_position - starting position of the bundle
+# * stop_position - ending position of the bundle.
+#
+# Type for start and stop positions are specific to the bounded source and must
+# be consistent throughout.
+SourceBundle = namedtuple(
+    'SourceBundle',
+    'weight source start_position stop_position')
+
+
+class BoundedSource(object):
+  """A Dataflow source that reads a finite amount of input records.
+
+  This class defines following operations which can be used to read the source
+  efficiently.
+  * Size estimation - method ``estimate_size()`` may return an accurate
+    estimation in bytes for the size of the source.
+  * Splitting into bundles of a given size - method ``split()`` can be used to
+    split the source into a set of sub-sources (bundles) based on a desired
+    bundle size.
+  * Getting a RangeTracker - method ``get_range_tracker() should return a
+    ``RangeTracker`` object for a given position range for the position type
+    of the records returned by the source.
+  * Reading the data - method ``read()`` can be used to read data from the
+    source while respecting the boundaries defined by a given
+    ``RangeTracker``.
+  """
+
+  def estimate_size(self):
+    """Estimates the size of source in bytes.
+
+    An estimate of the total size (in bytes) of the data that would be read
+    from this source. This estimate is in terms of external storage size,
+    before performing decompression or other processing.
+
+    Returns:
+      estimated size of the source if the size can be determined, ``None``
+      otherwise.
+    """
+    raise NotImplementedError
+
+  def split(self, desired_bundle_size, start_position=None, stop_position=None):
+    """Splits the source into a set of bundles.
+
+    Bundles should be approximately of size ``desired_bundle_size`` bytes.
+
+    Args:
+      desired_bundle_size: the desired size (in bytes) of the bundles returned.
+      start_position: if specified the given position must be used as the
+                      starting position of the first bundle.
+      stop_position: if specified the given position must be used as the ending
+                     position of the last bundle.
+    Returns:
+      an iterator of objects of type 'SourceBundle' that gives information about
+      the generated bundles.
+    """
+    raise NotImplementedError
+
+  def get_range_tracker(self, start_position, stop_position):
+    """Returns a RangeTracker for a given position range.
+
+    Framework may invoke ``read()`` method with the RangeTracker object returned
+    here to read data from the source.
+    Args:
+      start_position: starting position of the range.
+      stop_position:  ending position of the range.
+    Returns:
+      a ``RangeTracker`` for the given position range.
+    """
+    raise NotImplementedError
+
+  def read(self, range_tracker):
+    """Returns an iterator that reads data from the source.
+
+    The returned set of data must respect the boundaries defined by the given
+    ``RangeTracker`` object. For example:
+      * Returned set of data must be for the range
+        ``[range_tracker.start_position, range_tracker.stop_position)``. Note
+        that a source may decide to return records that start after
+        ``range_tracker.stop_position``. See documentation in class
+        ``RangeTracker`` for more details. Also, note that framework might
+        invoke ``range_tracker.try_split()`` to perform dynamic split
+        operations. range_tracker.stop_position may be updated
+        dynamically due to successful dynamic split operations.
+      * Method ``range_tracker.try_split()`` must be invoked for every record
+        that starts at a split point.
+      * Method ``range_tracker.record_current_position()`` may be invoked for
+        records that do not start at split points.
+    Args:
+      range_tracker: a ``RangeTracker`` whose boundaries must be respected
+                     when reading data from the source. If 'None' all records
+                     represented by the current source should be read.
+    Returns:
+      an iterator of data read by the source.
+    """
+    raise NotImplementedError
 
-  **Usage of the RangeTracker class hierarchy**
+  def default_output_coder(self):
+    """Coder that should be used for the records returned by the source."""
+    return PickleCoder()
 
-  The ``RangeTracker`` class should not be used per se---all users should use
-  its subclasses directly. We declare it here because all subclasses have
-  roughly the same interface and the same properties, to centralize the
-  documentation.
 
-  Currently we provide one implementation: ``iobase.OffsetRangeTracker``.
+class RangeTracker(object):
+  """A thread safe object used by Dataflow source framework.
+
+  A Dataflow source is defined using a ''BoundedSource'' and a ''RangeTracker''
+  pair. A ''RangeTracker'' is used by Dataflow source framework to perform
+  dynamic work rebalancing of position-based sources.
 
   **Position-based sources**
 
@@ -421,67 +532,110 @@ def stop_position(self):
     """Returns the ending position of the current range, exclusive."""
     raise NotImplementedError
 
-  def try_return_record_at(self, is_at_split_point, record_start):
-    """Atomically determines if a record at the given position can be returned.
+  def try_claim(self, position):  # pylint: disable=unused-argument
+    """Atomically determines if a record at a split point is within the range.
 
-    Additionally, Updates the internal state of the ``RangeTracker``.
+    This method should be called **if and only if** the record is at a split
+    point. This method may modify the internal state of the ``RangeTracker`` by
+    updating the last-consumed position to ``position``.
 
-    In particular:
+    ** Thread safety **
 
-    * If ``is_at_split_point`` is ``True``, and ``record_start`` is outside the
-      current range, returns ``False``;
-    * Otherwise, updates the last-consumed position to ``record_start`` and
-      returns ``True``.
+    This method along with several other methods of this class may be invoked by
+    multiple threads, hence must be made thread-safe, e.g. by using a single
+    lock object.
 
-    This method MUST be called on all split point records. It may be called on
-    every record.
+    Args:
+      position: starting position of a record being read by a source.
 
-    Method ``try_return_record_at`` and method ``try_split_at_position`` will be
-    accessed by different threads and implementor must ensure that only one of
-    these methods is executed at a given time.
+    Returns:
+      ``True``, if the given position falls within the current range, returns
+      ``False`` otherwise.
+    """
+    raise NotImplementedError
+
+  def set_current_position(self, position):
+    """Updates the last-consumed position to the given position.
+
+    A source may invoke this method for records that do not start at split
+    points. This may modify the internal state of the ``RangeTracker``. If the
+    record starts at a split point, method ``try_claim()`` **must** be invoked
+    instead of this method.
 
     Args:
-      is_at_split_point: ``True`` if record is at a split point, ``False``
-      otherwise.
+      position: starting position of a record being read by a source.
+    """
+    raise NotImplementedError
+
+  def position_at_fraction(self, fraction):
+    """Returns the position at the given fraction.
+
+    Given a fraction within the range [0.0, 1.0) this method will return the
+    position at the given fraction compared the the position range
+    [self.start_position, self.stop_position).
+
+    ** Thread safety **
 
-      record_start: starting position of the record.
+    This method along with several other methods of this class may be invoked by
+    multiple threads, hence must be made thread-safe, e.g. by using a single
+    lock object.
+
+    Args:
+      fraction: a float value within the range [0.0, 1.0).
+    Returns:
+      a position within the range [self.start_position, self.stop_position).
     """
     raise NotImplementedError
 
-  def try_split_at_position(self, split_position):
+  def try_split(self, position):
     """Atomically splits the current range.
 
-    Splits the current range '[get_start_position(), get_stop_position())'
-    into a "primary" part '[get_start_position(), split_position())' and a
-    "residual" part '[split_position(), get_stop_position())', assuming the
+    Determines a position to split the current range, split_position, based on
+    the given position. In most cases split_position and position will be the
+    same.
+
+    Splits the current range '[self.start_position, self.stop_position)'
+    into a "primary" part '[self.start_position, split_position)' and a
+    "residual" part '[split_position, self.stop_position)', assuming the
     current last-consumed position is within
-    '[get_start_position(), split_position())' (i.e., 'split_position()'
-    has not been consumed yet).
+    '[self.start_position, split_position)' (i.e., split_position has not been
+    consumed yet).
+
+    If successful, updates the current range to be the primary and returns a
+    tuple (split_position, split_fraction). split_fraction should be the
+    fraction of size of range '[self.start_position, split_position)' compared
+    to the original (before split) range
+    '[self.start_position, self.stop_position)'.
 
-    Updates the current range to be the primary and returns ``True``. This
-    means that all further calls on the current object will interpret their
-    arguments relative to the primary range.
+    If the split_position has already been consumed, returns ``None``.
 
-    If the split position has already been consumed, or if no
-    ``try_return_record_at`` call was made yet, returns ``False``. The
-    second condition is to prevent dynamic splitting during reader start-up.
+    ** Thread safety **
 
-    Method ``try_return_record_at`` and method ``try_split_at_position`` will be
-    accessed by different threads and implementor must ensure that only one of
-    these methods is executed at a given time.
+    This method along with several other methods of this class may be invoked by
+    multiple threads, hence must be made thread-safe, e.g. by using a single
+    lock object.
 
     Args:
-      split_position: an instance of ReaderPosition that gives the position
-      where the current range should be split at.
+      position: suggested position where the current range should try to
+                be split at.
+    Returns:
+      a tuple containing the split position and split fraction.
     """
     raise NotImplementedError
 
   def fraction_consumed(self):
     """Returns the approximate fraction of consumed positions in the source.
 
-    Returns the approximate fraction of positions that have been consumed by
-    successful 'try_return_record_at()' calls, or 0.0 if no such calls have
-    happened.
+    ** Thread safety **
+
+    This method along with several other methods of this class may be invoked by
+    multiple threads, hence must be made thread-safe, e.g. by using a single
+    lock object.
+
+    Returns:
+      the approximate fraction of positions that have been consumed by
+      successful 'try_split()' and  'report_current_position()'  calls, or
+      0.0 if no such calls have happened.
     """
     raise NotImplementedError