feat: add 'Rowset.add_row_range_with_prefix'

docs/snippets_table.py

@@ -620,7 +620,7 @@ def test_bigtable_create_update_delete_column_family():
     assert column_family_id not in column_families
 
 
-def test_bigtable_add_row_add_row_range_add_row_range_from_keys():
+def test_bigtable_add_row_add_row_range_add_row_range_from_keys_add_row_range_by_prefix_from_keys():
     row_keys = [
         b"row_key_1",
         b"row_key_2",
@@ -631,6 +631,7 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys():
         b"row_key_7",
         b"row_key_8",
         b"row_key_9",
+        b"sample_row_key_1",
     ]
 
     rows = []
@@ -692,6 +693,33 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys():
     found_row_keys = [row.row_key for row in read_rows]
     assert found_row_keys == expected_row_keys
 
+    # [START bigtable_row_range_by_prefix_from_keys]
+    from google.cloud.bigtable import Client
+    from google.cloud.bigtable.row_set import RowSet
+
+    client = Client(admin=True)
+    instance = client.instance(INSTANCE_ID)
+    table = instance.table(TABLE_ID)
+
+    row_set = RowSet()
+    row_set.add_row_range_by_prefix_from_keys("row")
+    # [END bigtable_row_range_by_prefix_from_keys]
+
+    read_rows = table.read_rows(row_set=row_set)
+    expected_row_keys = [
+        b"row_key_1",
+        b"row_key_2",
+        b"row_key_3",
+        b"row_key_4",
+        b"row_key_5",
+        b"row_key_6",
+        b"row_key_7",
+        b"row_key_8",
+        b"row_key_9",
+    ]
+    found_row_keys = [row.row_key for row in read_rows]
+    assert found_row_keys == expected_row_keys
+
     table.truncate(timeout=200)
 
 

google/cloud/bigtable/row_set.py

@@ -109,6 +109,24 @@ def add_row_range_from_keys(
         row_range = RowRange(start_key, end_key, start_inclusive, end_inclusive)
         self.row_ranges.append(row_range)
 
+    def add_row_range_by_prefix_from_keys(self, row_key_prefix):
+        """Add row range to row_ranges list that start with the row_key_prefix from the row keys
+
+        For example:
+
+        .. literalinclude:: snippets_table.py
+            :start-after: [START bigtable_row_range_by_prefix_from_keys]
+            :end-before: [END bigtable_row_range_by_prefix_from_keys]
+
+        :type row_key_prefix: str
+        :param row_key_prefix: To retrieve  all rows that start with this row key prefix.
+                            Prefix cannot be zero length."""
+
+        end_key = row_key_prefix[:-1] + chr(ord(row_key_prefix[-1]) + 1)
+        self.add_row_range_from_keys(
+            row_key_prefix.encode("utf-8"), end_key.encode("utf-8")
+        )
+
     def _update_message_request(self, message):
         """Add row keys and row range to given request message
 

tests/system.py

@@ -1013,6 +1013,38 @@ def test_yield_rows_with_row_set(self):
         found_row_keys = [row.row_key for row in read_rows]
         self.assertEqual(found_row_keys, expected_row_keys)
 
+    def test_add_row_range_by_prefix_from_keys(self):
+        row_keys = [
+            b"row_key_1",
+            b"row_key_2",
+            b"row_key_3",
+            b"row_key_4",
+            b"sample_row_key_1",
+            b"sample_row_key_2",
+        ]
+
+        rows = []
+        for row_key in row_keys:
+            row = self._table.row(row_key)
+            row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1)
+            rows.append(row)
+            self.rows_to_delete.append(row)
+        self._table.mutate_rows(rows)
+
+        row_set = RowSet()
+        row_set.add_row_range_by_prefix_from_keys("row")
+
+        read_rows = self._table.yield_rows(row_set=row_set)
+
+        expected_row_keys = [
+            b"row_key_1",
+            b"row_key_2",
+            b"row_key_3",
+            b"row_key_4",
+        ]
+        found_row_keys = [row.row_key for row in read_rows]
+        self.assertEqual(found_row_keys, expected_row_keys)
+
     def test_read_large_cell_limit(self):
         row = self._table.row(ROW_KEY)
         self.rows_to_delete.append(row)

tests/unit/test_row_set.py

@@ -167,6 +167,11 @@ def test_add_row_range_from_keys(self):
         )
         self.assertEqual(row_set.row_ranges[0].end_key, b"row_key9")
 
+    def test_add_row_range_by_prefix_from_keys(self):
+        row_set = self._make_one()
+        row_set.add_row_range_by_prefix_from_keys("row")
+        self.assertEqual(row_set.row_ranges[0].end_key, b"rox")
+
     def test__update_message_request(self):
         row_set = self._make_one()
         table_name = "table_name"