BigTable: Adding a row generator on a table. (#4679)

Allows iteration over the rows in a table instead of reading the rows into an internal dictionary first. As soon as a row has been validated, it is available in the iterator.

Author: aneepct

bigtable/google/cloud/bigtable/row_data.py

@@ -191,15 +191,64 @@ class PartialRowsData(object):
     :param response_iterator: A streaming iterator returned from a
                               ``ReadRows`` request.
     """
-    START = "Start"                         # No responses yet processed.
-    NEW_ROW = "New row"                     # No cells yet complete for row
-    ROW_IN_PROGRESS = "Row in progress"     # Some cells complete for row
-    CELL_IN_PROGRESS = "Cell in progress"   # Incomplete cell for row
+
+    START = 'Start'                         # No responses yet processed.
+    NEW_ROW = 'New row'                     # No cells yet complete for row
+    ROW_IN_PROGRESS = 'Row in progress'     # Some cells complete for row
+    CELL_IN_PROGRESS = 'Cell in progress'   # Incomplete cell for row
 
     def __init__(self, response_iterator):
         self._response_iterator = response_iterator
+        self._generator = YieldRowsData(response_iterator)
+
         # Fully-processed rows, keyed by `row_key`
-        self._rows = {}
+        self.rows = {}
+
+    def __eq__(self, other):
+        if not isinstance(other, self.__class__):
+            return NotImplemented
+        return other._response_iterator == self._response_iterator
+
+    def __ne__(self, other):
+        return not self == other
+
+    @property
+    def state(self):
+        """State machine state.
+
+        :rtype: str
+        :returns:  name of state corresponding to currrent row / chunk
+                   processing.
+        """
+        return self._generator.state
+
+    def consume_all(self, max_loops=None):
+        """Consume the streamed responses until there are no more.
+
+        :type max_loops: int
+        :param max_loops: (Optional) Maximum number of times to try to consume
+                          an additional ``ReadRowsResponse``. You can use this
+                          to avoid long wait times.
+        """
+        for row in self._generator.read_rows():
+            self.rows[row.row_key] = row
+
+
+class YieldRowsData(object):
+    """Convenience wrapper for consuming a ``ReadRows`` streaming response.
+
+    :type response_iterator: :class:`~google.cloud.exceptions.GrpcRendezvous`
+    :param response_iterator: A streaming iterator returned from a
+                              ``ReadRows`` request.
+    """
+
+    START = 'Start'  # No responses yet processed.
+    NEW_ROW = 'New row'  # No cells yet complete for row
+    ROW_IN_PROGRESS = 'Row in progress'  # Some cells complete for row
+    CELL_IN_PROGRESS = 'Cell in progress'  # Incomplete cell for row
+
+    def __init__(self, response_iterator):
+        self._response_iterator = response_iterator
         # Counter for responses pulled from iterator
         self._counter = 0
         # Maybe cached from previous response
@@ -213,14 +262,6 @@ def __init__(self, response_iterator):
         # Last complete cell, unset until first completion, after new row
         self._previous_cell = None
 
-    def __eq__(self, other):
-        if not isinstance(other, self.__class__):
-            return NotImplemented
-        return other._response_iterator == self._response_iterator
-
-    def __ne__(self, other):
-        return not self == other
-
     @property
     def state(self):
         """State machine state.
@@ -241,96 +282,75 @@ def state(self):
             return self.ROW_IN_PROGRESS
         return self.NEW_ROW  # row added, no chunk yet processed
 
-    @property
-    def rows(self):
-        """Property returning all rows accumulated from the stream.
-
-        :rtype: dict
-        :returns: row_key -> :class:`PartialRowData`.
-        """
-        # NOTE: To avoid duplicating large objects, this is just the
-        #       mutable private data.
-        return self._rows
-
     def cancel(self):
         """Cancels the iterator, closing the stream."""
         self._response_iterator.cancel()
 
-    def consume_next(self):
-        """Consume the next ``ReadRowsResponse`` from the stream.
+    def read_rows(self):
+        """Consume the ``ReadRowsResponse's`` from the stream.
+        Read the rows and yield each to the reader
 
         Parse the response and its chunks into a new/existing row in
         :attr:`_rows`. Rows are returned in order by row key.
         """
-        response = six.next(self._response_iterator)
-        self._counter += 1
+        while True:
+            try:
+                response = six.next(self._response_iterator)
+            except StopIteration:
+                break
 
-        if self._last_scanned_row_key is None:  # first response
-            if response.last_scanned_row_key:
-                raise InvalidReadRowsResponse()
+            self._counter += 1
 
-        self._last_scanned_row_key = response.last_scanned_row_key
+            if self._last_scanned_row_key is None:  # first response
+                if response.last_scanned_row_key:
+                    raise InvalidReadRowsResponse()
 
-        row = self._row
-        cell = self._cell
+            self._last_scanned_row_key = response.last_scanned_row_key
 
-        for chunk in response.chunks:
+            row = self._row
+            cell = self._cell
 
-            self._validate_chunk(chunk)
+            for chunk in response.chunks:
 
-            if chunk.reset_row:
-                row = self._row = None
-                cell = self._cell = self._previous_cell = None
-                continue
+                self._validate_chunk(chunk)
 
-            if row is None:
-                row = self._row = PartialRowData(chunk.row_key)
+                if chunk.reset_row:
+                    row = self._row = None
+                    cell = self._cell = self._previous_cell = None
+                    continue
 
-            if cell is None:
-                qualifier = None
-                if chunk.HasField('qualifier'):
-                    qualifier = chunk.qualifier.value
+                if row is None:
+                    row = self._row = PartialRowData(chunk.row_key)
 
-                cell = self._cell = PartialCellData(
-                    chunk.row_key,
-                    chunk.family_name.value,
-                    qualifier,
-                    chunk.timestamp_micros,
-                    chunk.labels,
-                    chunk.value)
-                self._copy_from_previous(cell)
-            else:
-                cell.append_value(chunk.value)
+                if cell is None:
+                    qualifier = None
+                    if chunk.HasField('qualifier'):
+                        qualifier = chunk.qualifier.value
 
-            if chunk.commit_row:
-                self._save_current_row()
-                row = cell = None
-                continue
+                    cell = self._cell = PartialCellData(
+                        chunk.row_key,
+                        chunk.family_name.value,
+                        qualifier,
+                        chunk.timestamp_micros,
+                        chunk.labels,
+                        chunk.value)
+                    self._copy_from_previous(cell)
+                else:
+                    cell.append_value(chunk.value)
 
-            if chunk.value_size == 0:
-                self._save_current_cell()
-                cell = None
+                if chunk.commit_row:
+                    self._save_current_cell()
 
-    def consume_all(self, max_loops=None):
-        """Consume the streamed responses until there are no more.
+                    yield self._row
 
-        This simply calls :meth:`consume_next` until there are no
-        more to consume.
+                    self._row, self._previous_row = None, self._row
+                    self._previous_cell = None
+                    row = cell = None
+                    continue
 
-        :type max_loops: int
-        :param max_loops: (Optional) Maximum number of times to try to consume
-                          an additional ``ReadRowsResponse``. You can use this
-                          to avoid long wait times.
-        """
-        curr_loop = 0
-        if max_loops is None:
-            max_loops = float('inf')
-        while curr_loop < max_loops:
-            curr_loop += 1
-            try:
-                self.consume_next()
-            except StopIteration:
-                break
+                if chunk.value_size == 0:
+                    self._save_current_cell()
+                    cell = None
 
     @staticmethod
     def _validate_chunk_status(chunk):
@@ -433,14 +453,6 @@ def _copy_from_previous(self, cell):
             if cell.qualifier is None:
                 cell.qualifier = previous.qualifier
 
-    def _save_current_row(self):
-        """Helper for :meth:`consume_next`."""
-        if self._cell:
-            self._save_current_cell()
-        self._rows[self._row.row_key] = self._row
-        self._row, self._previous_row = None, self._row
-        self._previous_cell = None
-
 
 def _raise_if(predicate, *args):
     """Helper for validation methods."""

bigtable/google/cloud/bigtable/table.py

@@ -31,6 +31,7 @@
 from google.cloud.bigtable.row import ConditionalRow
 from google.cloud.bigtable.row import DirectRow
 from google.cloud.bigtable.row_data import PartialRowsData
+from google.cloud.bigtable.row_data import YieldRowsData
 from grpc import StatusCode
 
 
@@ -315,6 +316,43 @@ def read_rows(self, start_key=None, end_key=None, limit=None,
         # We expect an iterator of `data_messages_v2_pb2.ReadRowsResponse`
         return PartialRowsData(response_iterator)
 
+    def yield_rows(self, start_key=None, end_key=None, limit=None,
+                   filter_=None):
+        """Read rows from this table.
+
+        :type start_key: bytes
+        :param start_key: (Optional) The beginning of a range of row keys to
+                          read from. The range will include ``start_key``. If
+                          left empty, will be interpreted as the empty string.
+
+        :type end_key: bytes
+        :param end_key: (Optional) The end of a range of row keys to read from.
+                        The range will not include ``end_key``. If left empty,
+                        will be interpreted as an infinite string.
+
+        :type limit: int
+        :param limit: (Optional) The read will terminate after committing to N
+                      rows' worth of results. The default (zero) is to return
+                      all results.
+
+        :type filter_: :class:`.RowFilter`
+        :param filter_: (Optional) The filter to apply to the contents of the
+                        specified row(s). If unset, reads every column in
+                        each row.
+
+        :rtype: :class:`.PartialRowData`
+        :returns: A :class:`.PartialRowData` for each row returned
+        """
+        request_pb = _create_row_request(
+            self.name, start_key=start_key, end_key=end_key, filter_=filter_,
+            limit=limit)
+        client = self._instance._client
+        response_iterator = client._data_stub.ReadRows(request_pb)
+        # We expect an iterator of `data_messages_v2_pb2.ReadRowsResponse`
+        generator = YieldRowsData(response_iterator)
+        for row in generator.read_rows():
+            yield row
+
     def mutate_rows(self, rows, retry=DEFAULT_RETRY):
         """Mutates multiple rows in bulk.