perf: improve row merging (#619)

The underlying GAPIC client uses protoplus for all requests and responses. However the underlying protos for ReadRowsResponse are never exposed to end users directly: the underlying chunks get merged into logic rows. The readability benefits provided by protoplus for ReadRows do not justify the costs. This change unwraps the protoplus messages and uses the raw protobuff message as input for row merging. This improves row merging performance by 10x. For 10k rows, each with 100 cells where each cell is 100 bytes and in groups of 100 rows per ReadRowsResponse, cProfile showed a 10x improvement:

old:          124266037 function calls in 68.208 seconds
new:          13042837 function calls in 7.787 seconds

There are still a few more low hanging fruits to optimize performance and those will come in follow up PRs

Author: igorbernstein2

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):