perf: improve row merging

google/cloud/bigtable/row_data.py

@@ -474,7 +474,11 @@ def _read_next(self):
 
     def _read_next_response(self):
         """Helper for :meth:`__iter__`."""
-        return self.retry(self._read_next, on_error=self._on_error)()
+        resp_protoplus = self.retry(self._read_next, on_error=self._on_error)()
+        # unwrap the underlying protobuf, there is a significant amount of
+        # overhead that protoplus imposes for very little gain. The protos
+        # are not user visible, so we just use the raw protos for merging.
+        return data_messages_v2_pb2.ReadRowsResponse.pb(resp_protoplus)
 
     def __iter__(self):
         """Consume the ``ReadRowsResponse`` s from the stream.
@@ -543,11 +547,12 @@ def _process_chunk(self, chunk):
     def _update_cell(self, chunk):
         if self._cell is None:
             qualifier = None
-            if "qualifier" in chunk:
-                qualifier = chunk.qualifier
+            if chunk.HasField("qualifier"):
+                qualifier = chunk.qualifier.value
+
             family = None
-            if "family_name" in chunk:
-                family = chunk.family_name
+            if chunk.HasField("family_name"):
+                family = chunk.family_name.value
 
             self._cell = PartialCellData(
                 chunk.row_key,
@@ -577,8 +582,8 @@ def _validate_chunk_reset_row(self, chunk):
 
         # No reset with other keys
         _raise_if(chunk.row_key)
-        _raise_if("family_name" in chunk)
-        _raise_if("qualifier" in chunk)
+        _raise_if(chunk.HasField("family_name"))
+        _raise_if(chunk.HasField("qualifier"))
         _raise_if(chunk.timestamp_micros)
         _raise_if(chunk.labels)
         _raise_if(chunk.value_size)

tests/unit/test_row_data.py

@@ -637,15 +637,15 @@ def test_partial_rows_data__copy_from_previous_filled():
 
 def test_partial_rows_data_valid_last_scanned_row_key_on_start():
     client = _Client()
-    response = _ReadRowsResponseV2(chunks=(), last_scanned_row_key="2.AFTER")
+    response = _ReadRowsResponseV2([], last_scanned_row_key=b"2.AFTER")
     iterator = _MockCancellableIterator(response)
     client._data_stub = mock.MagicMock()
     client._data_stub.read_rows.side_effect = [iterator]
     request = object()
     yrd = _make_partial_rows_data(client._data_stub.read_rows, request)
-    yrd.last_scanned_row_key = "1.BEFORE"
+    yrd.last_scanned_row_key = b"1.BEFORE"
     _partial_rows_data_consume_all(yrd)
-    assert yrd.last_scanned_row_key == "2.AFTER"
+    assert yrd.last_scanned_row_key == b"2.AFTER"
 
 
 def test_partial_rows_data_invalid_empty_chunk():
@@ -666,6 +666,7 @@ def test_partial_rows_data_invalid_empty_chunk():
 
 def test_partial_rows_data_state_cell_in_progress():
     from google.cloud.bigtable_v2.services.bigtable import BigtableClient
+    from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2
 
     LABELS = ["L1", "L2"]
 
@@ -682,6 +683,9 @@ def test_partial_rows_data_state_cell_in_progress():
         value=VALUE,
         labels=LABELS,
     )
+    # _update_cell expects to be called after the protoplus wrapper has been
+    # shucked
+    chunk = messages_v2_pb2.ReadRowsResponse.CellChunk.pb(chunk)
     yrd._update_cell(chunk)
 
     more_cell_data = _ReadRowsResponseCellChunkPB(value=VALUE)
@@ -1455,10 +1459,12 @@ def __init__(self, **kw):
         self.__dict__.update(kw)
 
 
-class _ReadRowsResponseV2(object):
-    def __init__(self, chunks, last_scanned_row_key=""):
-        self.chunks = chunks
-        self.last_scanned_row_key = last_scanned_row_key
+def _ReadRowsResponseV2(chunks, last_scanned_row_key=b""):
+    from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2
+
+    return messages_v2_pb2.ReadRowsResponse(
+        chunks=chunks, last_scanned_row_key=last_scanned_row_key
+    )
 
 
 def _generate_cell_chunks(chunk_text_pbs):

tests/unit/test_table.py

@@ -2206,10 +2206,12 @@ def next(self):
     __next__ = next
 
 
-class _ReadRowsResponseV2(object):
-    def __init__(self, chunks, last_scanned_row_key=""):
-        self.chunks = chunks
-        self.last_scanned_row_key = last_scanned_row_key
+def _ReadRowsResponseV2(chunks, last_scanned_row_key=b""):
+    from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2
+
+    return messages_v2_pb2.ReadRowsResponse(
+        chunks=chunks, last_scanned_row_key=last_scanned_row_key
+    )
 
 
 def _TablePB(*args, **kw):