Repository: googleapis/google-cloud-python

Issue:
BigTable: Memory usage abnormally increasing when using table.read_rows method
Hi,

While using the python google-cloud-bigtable to read a large collection of data (10M to 90M rows of at least 4MB to 5MB per row), I noticed a strange behavior: The used RAM was at lease twice the size of the data I was reading from the table.

Here is the code I did to reproduce :

```python
from google.cloud import bigtable

instance_name = "the-instance"
table_name = "the-table"

client = bigtable.Client(admin=False)
instance = client.instance(instance_name)
table = instance.table(table_name)

last_readall_key = None
rows = table.read_rows(start_key=last_readall_key)

while True:
    try:
        rows.consume_next()
    except StopIteration as e:
        break
    except Exception as e:
        print 'Got exception of type %s : %s' % (type(e), str(e)) # mostly grpc._channel._Rendezvous exceptions are thrown here

        table = instance.table(table_name)
        rows = table.read_rows(start_key=last_readall_key)

        rows.consume_next()

    for key in rows.rows:
        last_readall_key = key

    rows._rows.clear() #this was an attempt to fix, but it doesn't have any effect
```

When running the above code, my instance RAM fills itself until full.

Another weird thing, is that the problem does not happen on my MacOS Sierra 10.12.6 machine. I only witness the problem on Linux machines.

I saw this problem on Linux Ubuntu 14.04 LTS and Linux Ubuntu 10.04.5 LTS. I did not try with other Linux distributions or others versions of Ubuntu though.

Am I doing something wrong ? Does anyone knows what may cause this problem ?
Would anyone have any trick to avoid this problem ?

Find below the pip requirements to reproduce :

```
cachetools==2.0.1
certifi==2017.11.5
chardet==3.0.4
dill==0.2.7.1
enum34==1.1.6
future==0.16.0
futures==3.2.0
google-api-core==0.1.2
google-auth==1.2.1
google-cloud-bigtable==0.28.1
google-cloud-core==0.28.0
google-gax==0.15.16
googleapis-common-protos==1.5.3
grpcio==1.7.3
idna==2.6
ply==3.8
protobuf==3.5.0.post1
pyasn1==0.4.2
pyasn1-modules==0.2.1
pytz==2017.3
requests==2.18.4
rsa==3.4.2
six==1.11.0
urllib3==1.22
```
(these dependancies are installed automatically when using `pip install google-cloud-bigtable` )

Thanks for your time.


Discussion:
Hi, does anyone have an idea of how to solve this ?
@chemelnucfin this is much more of a bug than a question.

@Rafff Could you give an idea of how much data is in the table?
FYI, we're working on a yield_row implementation (https://github.com/zakons/google-cloud-python/tree/feature/row_iterator) that gets rid of the use of a map to track returned rows.  In other languages, we use iterators that return rows rather than using an intermediary map.  That type of iterator seems more like how a user thinks about interacting with a stream of results from Cloud Bigtable.
@dhermes the table is huge, it contains almost 90M rows where each row contains 1 cell of 4000Bytes.

@sduskis this is awesome ! That's exactly what I was trying to achieve, I was using the read_rows to yield row by row to my application. When do you think this could be released ?

Thanks for your time !

Resolving patch:
diff --git a/bigtable/google/cloud/bigtable/row_data.py b/bigtable/google/cloud/bigtable/row_data.py
--- a/bigtable/google/cloud/bigtable/row_data.py
+++ b/bigtable/google/cloud/bigtable/row_data.py
@@ -187,15 +187,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
@@ -209,14 +258,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.
@@ -237,96 +278,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):
@@ -429,14 +449,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."""
diff --git a/bigtable/google/cloud/bigtable/table.py b/bigtable/google/cloud/bigtable/table.py
--- a/bigtable/google/cloud/bigtable/table.py
+++ b/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.
 


Test patch:
diff --git a/bigtable/tests/unit/test_row_data.py b/bigtable/tests/unit/test_row_data.py
--- a/bigtable/tests/unit/test_row_data.py
+++ b/bigtable/tests/unit/test_row_data.py
@@ -14,8 +14,8 @@
 
 
 import unittest
-
-import mock
+import timeit
+import csv
 
 
 class TestCell(unittest.TestCase):
@@ -196,22 +196,6 @@ def _get_target_class():
 
         return PartialRowsData
 
-    def _getDoNothingClass(self):
-        klass = self._get_target_class()
-
-        class FakePartialRowsData(klass):
-
-            def __init__(self, *args, **kwargs):
-                super(FakePartialRowsData, self).__init__(*args, **kwargs)
-                self._consumed = []
-
-            def consume_next(self):
-                value = self._response_iterator.next()
-                self._consumed.append(value)
-                return value
-
-        return FakePartialRowsData
-
     def _make_one(self, *args, **kwargs):
         return self._get_target_class()(*args, **kwargs)
 
@@ -220,7 +204,7 @@ def test_constructor(self):
         partial_rows_data = self._make_one(response_iterator)
         self.assertIs(partial_rows_data._response_iterator,
                       response_iterator)
-        self.assertEqual(partial_rows_data._rows, {})
+        self.assertEqual(partial_rows_data.rows, {})
 
     def test___eq__(self):
         response_iterator = object()
@@ -247,59 +231,52 @@ def test___ne__(self):
         partial_rows_data2 = self._make_one(response_iterator2)
         self.assertNotEqual(partial_rows_data1, partial_rows_data2)
 
-    def test_state_start(self):
-        prd = self._make_one([])
-        self.assertEqual(prd.state, prd.START)
-
-    def test_state_new_row_w_row(self):
-        prd = self._make_one([])
-        prd._last_scanned_row_key = ''
-        prd._row = object()
-        self.assertEqual(prd.state, prd.NEW_ROW)
-
     def test_rows_getter(self):
         partial_rows_data = self._make_one(None)
-        partial_rows_data._rows = value = object()
+        partial_rows_data.rows = value = object()
         self.assertIs(partial_rows_data.rows, value)
 
-    def test_cancel(self):
-        response_iterator = _MockCancellableIterator()
-        partial_rows_data = self._make_one(response_iterator)
-        self.assertEqual(response_iterator.cancel_calls, 0)
-        partial_rows_data.cancel()
-        self.assertEqual(response_iterator.cancel_calls, 1)
 
-    # 'consume_nest' tested via 'TestPartialRowsData_JSON_acceptance_tests'
+class TestYieldRowsData(unittest.TestCase):
+    ROW_KEY = b'row-key'
+    FAMILY_NAME = u'family'
+    QUALIFIER = b'qualifier'
+    TIMESTAMP_MICROS = 100
+    VALUE = b'value'
+
+    @staticmethod
+    def _get_target_class():
+        from google.cloud.bigtable.row_data import YieldRowsData
+
+        return YieldRowsData
+
+    def _make_one(self, *args, **kwargs):
+        return self._get_target_class()(*args, **kwargs)
 
-    def test_consume_all(self):
-        klass = self._getDoNothingClass()
+    def test_state_start(self):
+        yrd = self._make_one([])
+        self.assertEqual(yrd.state, yrd.START)
 
-        value1, value2, value3 = object(), object(), object()
-        response_iterator = _MockCancellableIterator(value1, value2, value3)
-        partial_rows_data = klass(response_iterator)
-        self.assertEqual(partial_rows_data._consumed, [])
-        partial_rows_data.consume_all()
-        self.assertEqual(
-            partial_rows_data._consumed, [value1, value2, value3])
+    def test_state_new_row_w_row(self):
+        yrd = self._make_one([])
+        yrd._last_scanned_row_key = ''
+        yrd._row = object()
+        self.assertEqual(yrd.state, yrd.NEW_ROW)
 
-    def test_consume_all_with_max_loops(self):
-        klass = self._getDoNothingClass()
+    def test_cancel(self):
+        response_iterator = _MockCancellableIterator()
+        yield_rows_data = self._make_one(response_iterator)
+        self.assertEqual(response_iterator.cancel_calls, 0)
+        yield_rows_data.cancel()
+        self.assertEqual(response_iterator.cancel_calls, 1)
 
-        value1, value2, value3 = object(), object(), object()
-        response_iterator = _MockCancellableIterator(value1, value2, value3)
-        partial_rows_data = klass(response_iterator)
-        self.assertEqual(partial_rows_data._consumed, [])
-        partial_rows_data.consume_all(max_loops=1)
-        self.assertEqual(partial_rows_data._consumed, [value1])
-        # Make sure the iterator still has the remaining values.
-        self.assertEqual(
-            list(response_iterator.iter_values), [value2, value3])
+    # 'consume_next' tested via 'TestPartialRowsData_JSON_acceptance_tests'
 
     def test__copy_from_current_unset(self):
-        prd = self._make_one([])
+        yrd = self._make_one([])
         chunks = _generate_cell_chunks([''])
         chunk = chunks[0]
-        prd._copy_from_current(chunk)
+        yrd._copy_from_current(chunk)
         self.assertEqual(chunk.row_key, b'')
         self.assertEqual(chunk.family_name.value, u'')
         self.assertEqual(chunk.qualifier.value, b'')
@@ -312,8 +289,8 @@ def test__copy_from_current_blank(self):
         QUALIFIER = b'C'
         TIMESTAMP_MICROS = 100
         LABELS = ['L1', 'L2']
-        prd = self._make_one([])
-        prd._cell = _PartialCellData()
+        yrd = self._make_one([])
+        yrd._cell = _PartialCellData()
         chunks = _generate_cell_chunks([''])
         chunk = chunks[0]
         chunk.row_key = ROW_KEY
@@ -321,7 +298,7 @@ def test__copy_from_current_blank(self):
         chunk.qualifier.value = QUALIFIER
         chunk.timestamp_micros = TIMESTAMP_MICROS
         chunk.labels.extend(LABELS)
-        prd._copy_from_current(chunk)
+        yrd._copy_from_current(chunk)
         self.assertEqual(chunk.row_key, ROW_KEY)
         self.assertEqual(chunk.family_name.value, FAMILY_NAME)
         self.assertEqual(chunk.qualifier.value, QUALIFIER)
@@ -329,9 +306,9 @@ def test__copy_from_current_blank(self):
         self.assertEqual(chunk.labels, LABELS)
 
     def test__copy_from_previous_unset(self):
-        prd = self._make_one([])
+        yrd = self._make_one([])
         cell = _PartialCellData()
-        prd._copy_from_previous(cell)
+        yrd._copy_from_previous(cell)
         self.assertEqual(cell.row_key, '')
         self.assertEqual(cell.family_name, u'')
         self.assertIsNone(cell.qualifier)
@@ -344,7 +321,7 @@ def test__copy_from_previous_blank(self):
         QUALIFIER = b'C'
         TIMESTAMP_MICROS = 100
         LABELS = ['L1', 'L2']
-        prd = self._make_one([])
+        yrd = self._make_one([])
         cell = _PartialCellData(
             row_key=ROW_KEY,
             family_name=FAMILY_NAME,
@@ -352,8 +329,8 @@ def test__copy_from_previous_blank(self):
             timestamp_micros=TIMESTAMP_MICROS,
             labels=LABELS,
         )
-        prd._previous_cell = _PartialCellData()
-        prd._copy_from_previous(cell)
+        yrd._previous_cell = _PartialCellData()
+        yrd._copy_from_previous(cell)
         self.assertEqual(cell.row_key, ROW_KEY)
         self.assertEqual(cell.family_name, FAMILY_NAME)
         self.assertEqual(cell.qualifier, QUALIFIER)
@@ -366,8 +343,8 @@ def test__copy_from_previous_filled(self):
         QUALIFIER = b'C'
         TIMESTAMP_MICROS = 100
         LABELS = ['L1', 'L2']
-        prd = self._make_one([])
-        prd._previous_cell = _PartialCellData(
+        yrd = self._make_one([])
+        yrd._previous_cell = _PartialCellData(
             row_key=ROW_KEY,
             family_name=FAMILY_NAME,
             qualifier=QUALIFIER,
@@ -375,38 +352,30 @@ def test__copy_from_previous_filled(self):
             labels=LABELS,
         )
         cell = _PartialCellData()
-        prd._copy_from_previous(cell)
+        yrd._copy_from_previous(cell)
         self.assertEqual(cell.row_key, ROW_KEY)
         self.assertEqual(cell.family_name, FAMILY_NAME)
         self.assertEqual(cell.qualifier, QUALIFIER)
         self.assertEqual(cell.timestamp_micros, 0)
         self.assertEqual(cell.labels, [])
 
-    def test__save_row_no_cell(self):
-        ROW_KEY = 'RK'
-        prd = self._make_one([])
-        row = prd._row = mock.Mock(row_key=ROW_KEY, spec=['row_key'])
-        prd._cell = None
-        prd._save_current_row()
-        self.assertIs(prd._rows[ROW_KEY], row)
-
     def test_invalid_last_scanned_row_key_on_start(self):
         from google.cloud.bigtable.row_data import InvalidReadRowsResponse
 
         response = _ReadRowsResponseV2(chunks=(), last_scanned_row_key='ABC')
         iterator = _MockCancellableIterator(response)
-        prd = self._make_one(iterator)
+        yrd = self._make_one(iterator)
         with self.assertRaises(InvalidReadRowsResponse):
-            prd.consume_next()
+            self._consume_all(yrd)
 
     def test_valid_last_scanned_row_key_on_start(self):
         response = _ReadRowsResponseV2(
             chunks=(), last_scanned_row_key='AFTER')
         iterator = _MockCancellableIterator(response)
-        prd = self._make_one(iterator)
-        prd._last_scanned_row_key = 'BEFORE'
-        prd.consume_next()
-        self.assertEqual(prd._last_scanned_row_key, 'AFTER')
+        yrd = self._make_one(iterator)
+        yrd._last_scanned_row_key = 'BEFORE'
+        self._consume_all(yrd)
+        self.assertEqual(yrd._last_scanned_row_key, 'AFTER')
 
     def test_invalid_empty_chunk(self):
         from google.cloud.bigtable.row_data import InvalidChunk
@@ -414,9 +383,36 @@ def test_invalid_empty_chunk(self):
         chunks = _generate_cell_chunks([''])
         response = _ReadRowsResponseV2(chunks)
         iterator = _MockCancellableIterator(response)
-        prd = self._make_one(iterator)
+        yrd = self._make_one(iterator)
         with self.assertRaises(InvalidChunk):
-            prd.consume_next()
+            self._consume_all(yrd)
+
+    def test_yield_rows_data(self):
+        from google.cloud.bigtable.row_data import YieldRowsData
+
+        chunk = _ReadRowsResponseCellChunkPB(
+            row_key=self.ROW_KEY,
+            family_name=self.FAMILY_NAME,
+            qualifier=self.QUALIFIER,
+            timestamp_micros=self.TIMESTAMP_MICROS,
+            value=self.VALUE,
+            commit_row=True,
+        )
+        chunks = [chunk]
+
+        response = _ReadRowsResponseV2(chunks)
+        iterator = _MockCancellableIterator(response)
+        yrd = YieldRowsData(iterator)
+
+        rows = []
+        for row in yrd.read_rows():
+            rows.append(row)
+        result = rows[0]
+
+        self.assertEqual(result.row_key, self.ROW_KEY)
+
+    def _consume_all(self, yrd):
+        return [row.row_key for row in yrd.read_rows()]
 
 
 class TestPartialRowsData_JSON_acceptance_tests(unittest.TestCase):
@@ -454,7 +450,7 @@ def _fail_during_consume(self, testcase_name):
         iterator = _MockCancellableIterator(response)
         prd = self._make_one(iterator)
         with self.assertRaises(InvalidChunk):
-            prd.consume_next()
+            prd.consume_all()
         expected_result = self._sort_flattend_cells(
             [result for result in results if not result['error']])
         flattened = self._sort_flattend_cells(_flatten_cells(prd))
@@ -510,7 +506,7 @@ def _incomplete_final_row(self, testcase_name):
         response = _ReadRowsResponseV2(chunks)
         iterator = _MockCancellableIterator(response)
         prd = self._make_one(iterator)
-        prd.consume_next()
+        prd.consume_all()
         self.assertEqual(prd.state, prd.ROW_IN_PROGRESS)
         expected_result = self._sort_flattend_cells(
             [result for result in results if not result['error']])
@@ -532,7 +528,7 @@ def _match_results(self, testcase_name, expected_result=_marker):
         response = _ReadRowsResponseV2(chunks)
         iterator = _MockCancellableIterator(response)
         prd = self._make_one(iterator)
-        prd.consume_next()
+        prd.consume_all()
         flattened = self._sort_flattend_cells(_flatten_cells(prd))
         if expected_result is self._marker:
             expected_result = self._sort_flattend_cells(results)
@@ -727,3 +723,15 @@ def _parse_readrows_acceptance_tests(filename):
         chunks = _generate_cell_chunks(test['chunks'])
         results = test['results']
         yield name, chunks, results
+
+
+def _ReadRowsResponseCellChunkPB(*args, **kw):
+    from google.cloud.bigtable._generated import (
+        bigtable_pb2 as messages_v2_pb2)
+
+    family_name = kw.pop('family_name')
+    qualifier = kw.pop('qualifier')
+    message = messages_v2_pb2.ReadRowsResponse.CellChunk(*args, **kw)
+    message.family_name.value = family_name
+    message.qualifier.value = qualifier
+    return message
diff --git a/bigtable/tests/unit/test_table.py b/bigtable/tests/unit/test_table.py
--- a/bigtable/tests/unit/test_table.py
+++ b/bigtable/tests/unit/test_table.py
@@ -131,6 +131,7 @@ class TestTable(unittest.TestCase):
     QUALIFIER = b'qualifier'
     TIMESTAMP_MICROS = 100
     VALUE = b'value'
+    _json_tests = None
 
     @staticmethod
     def _get_target_class():
@@ -521,6 +522,37 @@ def mock_create_row_request(table_name, **kwargs):
         }
         self.assertEqual(mock_created, [(table.name, created_kwargs)])
 
+    def test_yield_rows(self):
+        from tests.unit._testing import _FakeStub
+
+        client = _Client()
+        instance = _Instance(self.INSTANCE_NAME, client=client)
+        table = self._make_one(self.TABLE_ID, instance)
+
+        # Create response_iterator
+        chunk = _ReadRowsResponseCellChunkPB(
+            row_key=self.ROW_KEY,
+            family_name=self.FAMILY_NAME,
+            qualifier=self.QUALIFIER,
+            timestamp_micros=self.TIMESTAMP_MICROS,
+            value=self.VALUE,
+            commit_row=True,
+        )
+        chunks = [chunk]
+
+        response = _ReadRowsResponseV2(chunks)
+        response_iterator = _MockCancellableIterator(response)
+
+        # Patch the stub used by the API method.
+        client._data_stub = _FakeStub(response_iterator)
+
+        rows = []
+        for row in table.yield_rows():
+            rows.append(row)
+        result = rows[0]
+
+        self.assertEqual(result.row_key, self.ROW_KEY)
+
     def test_sample_row_keys(self):
         from tests.unit._testing import _FakeStub
 
@@ -1184,3 +1216,24 @@ class _Instance(object):
     def __init__(self, name, client=None):
         self.name = name
         self._client = client
+
+
+class _MockCancellableIterator(object):
+
+    cancel_calls = 0
+
+    def __init__(self, *values):
+        self.iter_values = iter(values)
+
+    def next(self):
+        return next(self.iter_values)
+
+    def __next__(self):  # pragma: NO COVER Py3k
+        return self.next()
+
+
+class _ReadRowsResponseV2(object):
+
+    def __init__(self, chunks, last_scanned_row_key=''):
+        self.chunks = chunks
+        self.last_scanned_row_key = last_scanned_row_key