From d5969da8032682f0c571247c5ea0229b6cfdd420 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Thu, 4 Aug 2022 13:27:44 -0400 Subject: [PATCH 1/4] fix: improve row merging perf by 10x 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 --- google/cloud/bigtable/row_data.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/google/cloud/bigtable/row_data.py b/google/cloud/bigtable/row_data.py index 0c1565737..ab0358285 100644 --- a/google/cloud/bigtable/row_data.py +++ b/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) From 609cdddd557448805284411dec407fb350ddb7c5 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Thu, 4 Aug 2022 16:45:26 -0400 Subject: [PATCH 2/4] fix tests The previous approach of duck typing the protobuf messages to plain python objects no longer works as we need to shuck protoplus now --- tests/unit/test_row_data.py | 18 ++++++++++-------- tests/unit/test_table.py | 7 +++---- 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/tests/unit/test_row_data.py b/tests/unit/test_row_data.py index 9b329dc9f..30c8f61c0 100644 --- a/tests/unit/test_row_data.py +++ b/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,11 +1459,9 @@ 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): from google.protobuf.text_format import Merge diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py index 883f713d8..3aca6aaf2 100644 --- a/tests/unit/test_table.py +++ b/tests/unit/test_table.py @@ -2206,10 +2206,9 @@ 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): From f310b6be30a2bb6366918412cc9c7d7556437175 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Thu, 4 Aug 2022 16:54:42 -0400 Subject: [PATCH 3/4] lint --- tests/unit/test_row_data.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/test_row_data.py b/tests/unit/test_row_data.py index 30c8f61c0..7149de8b0 100644 --- a/tests/unit/test_row_data.py +++ b/tests/unit/test_row_data.py @@ -1463,6 +1463,7 @@ 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): from google.protobuf.text_format import Merge from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse From a4debf46d5726963ed449b1a4efb109b9624f2f0 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Thu, 4 Aug 2022 17:01:20 -0400 Subject: [PATCH 4/4] lint again --- tests/unit/test_row_data.py | 5 ++++- tests/unit/test_table.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_row_data.py b/tests/unit/test_row_data.py index 7149de8b0..94a90aa24 100644 --- a/tests/unit/test_row_data.py +++ b/tests/unit/test_row_data.py @@ -1461,7 +1461,10 @@ def __init__(self, **kw): 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) + + return messages_v2_pb2.ReadRowsResponse( + chunks=chunks, last_scanned_row_key=last_scanned_row_key + ) def _generate_cell_chunks(chunk_text_pbs): diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py index 3aca6aaf2..a89e02e8c 100644 --- a/tests/unit/test_table.py +++ b/tests/unit/test_table.py @@ -2208,7 +2208,10 @@ def next(self): 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) + + return messages_v2_pb2.ReadRowsResponse( + chunks=chunks, last_scanned_row_key=last_scanned_row_key + ) def _TablePB(*args, **kw):