googleapis · tseaver · Sep 21, 2018 · Sep 16, 2018 · Sep 16, 2018 · Sep 20, 2018
diff --git a/src/google/cloud/happybase/table.py b/src/google/cloud/happybase/table.py
@@ -35,6 +35,7 @@
 from google.cloud.bigtable.row_filters import TimestampRange
 from google.cloud.bigtable.row_filters import TimestampRangeFilter
 from google.cloud.bigtable.table import Table as _LowLevelTable
+from google.cloud.bigtable.row_set import RowSet
 
 from google.cloud.happybase.batch import _get_column_pairs
 from google.cloud.happybase.batch import _WAL_SENTINEL
@@ -235,7 +236,7 @@ def rows(self, rows, columns=None, timestamp=None,
         # versions == 1 since we only want the latest.
         filter_ = _filter_chain_helper(versions=1, timestamp=timestamp,
                                        filters=filters)
-        rows_generator = self._low_level_table.yield_rows(
+        rows_generator = self._low_level_table.read_rows(
             filter_=filter_)
         # NOTE: We could use max_loops = 1000 or some similar value to ensure
         #       that the stream isn't open too long.
@@ -381,9 +382,10 @@ def scan(self, row_start=None, row_stop=None, row_prefix=None,
         row_start, row_stop, filter_chain = _scan_filter_helper(
             row_start, row_stop, row_prefix, columns, timestamp, limit, kwargs)
 
-        rows_generator = self._low_level_table.yield_rows(
-            start_key=row_start, end_key=row_stop,
-            limit=limit, filter_=filter_chain)
+        row_set = _get_row_set_object(row_start, row_stop)
+
+        rows_generator = self._low_level_table.read_rows(
+            row_set=row_set, limit=limit, filter_=filter_chain)
 
         for rowdata in rows_generator:
             curr_row_data = rowdata
@@ -969,3 +971,12 @@ def _row_keys_filter_helper(row_keys):
         return filters[0]
     else:
         return RowFilterUnion(filters=filters)
+
+
+def _get_row_set_object(row_start, row_stop):
+    """Return a RowSet object for the given row_start and row_stop
+    """
+    row_set = RowSet()
+    row_set.add_row_range_from_keys(start_key=row_start,
+                                    end_key=row_stop)
+    return row_set
diff --git a/unit_tests/test_table.py b/unit_tests/test_table.py
@@ -619,8 +619,14 @@ def mock_filter_chain_helper(**kwargs):
             mock_filters.append(kwargs)
             return fake_filter
 
+        fake_row_set = object()
+
+        def mock_get_row_set_object(*args):  # pylint: disable=unused-argument
+            return fake_row_set
+
         with _Monkey(MUT, _filter_chain_helper=mock_filter_chain_helper,
-                     _columns_filter_helper=mock_columns_filter_helper):
+                     _columns_filter_helper=mock_columns_filter_helper,
+                     _get_row_set_object=mock_get_row_set_object):
             result = table.scan(row_start=row_start, row_stop=row_stop,
                                 row_prefix=row_prefix, columns=columns,
                                 filter=filter_, timestamp=timestamp,
@@ -636,12 +642,13 @@ def mock_filter_chain_helper(**kwargs):
         if row_prefix:
             row_start = row_prefix
             row_stop = MUT._string_successor(row_prefix)
+
         read_rows_kwargs = {
-            'end_key': row_stop,
+            'row_set': fake_row_set,
             'filter_': fake_filter,
             'limit': limit,
-            'start_key': row_start,
         }
+
         self.assertEqual(table._low_level_table.read_rows_calls, [
             (read_rows_args, read_rows_kwargs),
         ])
@@ -1427,6 +1434,22 @@ def test_many_rows(self):
         self.assertEqual(result, expected_result)
 
 
+class Test___get_row_set_object(unittest.TestCase):
+
+    def _callFUT(self, *args, **kwargs):
+        from google.cloud.happybase.table import _get_row_set_object
+        return _get_row_set_object(*args, **kwargs)
+
+    def test_row_set_object(self):
+        from google.cloud.bigtable.row_set import RowSet
+
+        start_key = b"row_key2"
+        end_key = b"row_key9"
+
+        row_set = self._callFUT(start_key, end_key)
+        self.assertIsInstance(row_set, RowSet)
+
+
 class _Connection(object):
 
     def __init__(self, instance):
@@ -1466,7 +1489,7 @@ def read_row(self, *args, **kwargs):
         self.read_row_calls.append((args, kwargs))
         return self.read_row_result
 
-    def yield_rows(self, *args, **kwargs):
+    def read_rows(self, *args, **kwargs):
         self.read_rows_calls.append((args, kwargs))
         rows_dict = self.read_rows_result.rows
         for row_key in sorted(rows_dict):