From f8d3b11818b2fa3cb3ce049c71882592727ab5c3 Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Thu, 15 Oct 2026 20:03:37 +0000 Subject: [PATCH 1/6] fix(bigtable): resolve O(n^2) bottleneck in ReadRowsQuery.add_range --- .../cloud/bigtable/data/read_rows_query.py | 22 ++++++++++++++++-- .../tests/unit/data/test_read_rows_query.py | 23 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py index 9fd42e2d871a..1038ce0bafe8 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py @@ -87,6 +87,17 @@ def __init__( init_dict["end_key_open"] = end_key self._pb = RowRangePB(**init_dict) + def _key(self) -> tuple[bytes | None, bool, bytes | None, bool]: + return ( + self.start_key, + self.start_is_inclusive, + self.end_key, + self.end_is_inclusive, + ) + + def __hash__(self) -> int: + return hash(self._key()) + @property def start_key(self) -> bytes | None: """ @@ -186,7 +197,7 @@ def __bool__(self) -> bool: def __eq__(self, other: Any) -> bool: if not isinstance(other, RowRange): return NotImplemented - return self._pb == other._pb + return self._key() == other._key() def __str__(self) -> str: """ @@ -250,6 +261,11 @@ def __init__( self._row_set = RowSetPB( row_keys=row_keys, row_ranges=[r._pb for r in row_ranges] ) + # Pre-initialize our signature set using hashable RowRange wrappers + self._range_signatures = set( + RowRange._from_pb(r) for r in self._row_set.row_ranges + ) + self.limit = limit or None self.filter = row_filter @@ -349,7 +365,9 @@ def add_range( Args: row_range: a range of row keys to add to this query """ - if row_range not in self.row_ranges: + # Utilize pure O(1) set-membership checking + if row_range not in self._range_signatures: + self._range_signatures.add(row_range) self._row_set.row_ranges.append(row_range._pb) def shard(self, shard_keys: RowKeySamples) -> ShardedQuery: diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py b/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py index e88ca4a9bb71..0758f17a55f2 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import time + import pytest TEST_ROWS = [ @@ -586,3 +588,24 @@ def test_empty_row_set(self): """Empty strings should be treated as keys inputs""" query = self._make_one(row_keys="") assert query.row_keys == [b""] + + def test_add_range_performance_regression(self): + """Assert O(1) range insertion scaling and prevent regression back to O(N^2).""" + from google.cloud.bigtable.data.read_rows_query import RowRange + + query = self._make_one() + ranges_count = 1000 + + start_time = time.perf_counter() + for i in range(ranges_count): + # Using 4-digit zero-padding ensures 'key_0009' is lexicographically smaller than 'key_0010' + r = RowRange(start_key=f"key_{i:04d}", end_key=f"key_{i + 1:04d}") + query.add_range(r) + end_time = time.perf_counter() + + elapsed = end_time - start_time + # Under O(N^2) on test runners, 1k iterations takes seconds. + # Under O(1) with tuple hashing, it takes <0.02s. + assert elapsed < 0.5, ( + f"O(N^2) Performance regression detected! Inserting {ranges_count} ranges took {elapsed:.2f}s" + ) From 679ebedf210f283bdd754aedbc3e7455d177ea54 Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Thu, 15 Oct 2026 20:09:00 +0000 Subject: [PATCH 2/6] address offline feedback --- .../google/cloud/bigtable/data/read_rows_query.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py index 1038ce0bafe8..84de74e31923 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py @@ -151,14 +151,14 @@ def _to_pb(self) -> RowRangePB: @classmethod def _from_pb(cls, data: RowRangePB) -> RowRange: """ - Creates a RowRange from a protobuf + Creates a RowRange from a protobuf without allocating temporary unused objects. Args: data (RowRangePB): The protobuf to convert Returns: RowRange: The converted RowRange """ - instance = cls() + instance = cls.__new__(cls) instance._pb = data return instance From 285b2c1f067c8c9d1a12058a93a977cd6fff99eb Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Thu, 15 Oct 2026 20:21:15 +0000 Subject: [PATCH 3/6] update comment --- .../tests/unit/data/test_read_rows_query.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py b/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py index 0758f17a55f2..a9b0dd9850bb 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py @@ -605,7 +605,8 @@ def test_add_range_performance_regression(self): elapsed = end_time - start_time # Under O(N^2) on test runners, 1k iterations takes seconds. - # Under O(1) with tuple hashing, it takes <0.02s. - assert elapsed < 0.5, ( + # Under O(1) with tuple hashing, it takes <0.10s. + # Double the value to avoid flaky tests + assert elapsed < 0.20, ( f"O(N^2) Performance regression detected! Inserting {ranges_count} ranges took {elapsed:.2f}s" ) From 01e759ad9fdcb2626b699bf41b5903b7ae3ed67d Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Thu, 15 Oct 2026 21:06:05 +0000 Subject: [PATCH 4/6] fix build; add test --- .../cloud/bigtable/data/read_rows_query.py | 2 +- .../tests/unit/data/test_read_rows_query.py | 23 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py index 84de74e31923..a15b51db7cb0 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py @@ -197,7 +197,7 @@ def __bool__(self) -> bool: def __eq__(self, other: Any) -> bool: if not isinstance(other, RowRange): return NotImplemented - return self._key() == other._key() + return self._pb == other._pb def __str__(self) -> str: """ diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py b/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py index a9b0dd9850bb..84d2aa577df8 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py @@ -610,3 +610,26 @@ def test_add_range_performance_regression(self): assert elapsed < 0.20, ( f"O(N^2) Performance regression detected! Inserting {ranges_count} ranges took {elapsed:.2f}s" ) + + def test_add_range_deduplication_behavior(self): + """Assert that identical ranges are deduplicated, but distinct ranges are kept.""" + from google.cloud.bigtable.data.read_rows_query import RowRange + + query = self._make_one() + + # Create two identical ranges + r1 = RowRange(start_key="a", end_key="b") + r2 = RowRange(start_key="a", end_key="b") + + # Create a distinct range + r3 = RowRange(start_key="c", end_key="d") + + # Add all three ranges + query.add_range(r1) + query.add_range(r2) # Duplicate, should be ignored + query.add_range(r3) # Distinct, should be kept + + # Assert that only r1 and r3 are in the final query + assert len(query.row_ranges) == 2 + assert query.row_ranges[0] == r1 + assert query.row_ranges[1] == r3 From d08cc34367f1031492b1c1afd1fd50c585308d03 Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Thu, 15 Oct 2026 21:18:34 +0000 Subject: [PATCH 5/6] address review feedback --- .../tests/unit/data/test_read_rows_query.py | 47 ++++++++++++------- 1 file changed, 29 insertions(+), 18 deletions(-) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py b/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py index 84d2aa577df8..66eaa26663ad 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py @@ -589,27 +589,38 @@ def test_empty_row_set(self): query = self._make_one(row_keys="") assert query.row_keys == [b""] - def test_add_range_performance_regression(self): - """Assert O(1) range insertion scaling and prevent regression back to O(N^2).""" + def test_add_range_does_not_materialize_row_ranges(self): + """ + Ensures add_range scales in O(1) time and does not trigger O(N) scans. + + The legacy implementation called the `row_ranges` property helper during + membership checks, which dynamically re-allocated and materialized brand + new Python RowRange wrapper objects for every element in the array on + each insertion. This led to an accidental O(N^2) memory and CPU leak. + + This test mocks out the `row_ranges` property on the class and asserts + that it is never accessed during batch additions, proving we bypass the + quadratic evaluation path entirely. + """ + from unittest import mock from google.cloud.bigtable.data.read_rows_query import RowRange query = self._make_one() - ranges_count = 1000 - - start_time = time.perf_counter() - for i in range(ranges_count): - # Using 4-digit zero-padding ensures 'key_0009' is lexicographically smaller than 'key_0010' - r = RowRange(start_key=f"key_{i:04d}", end_key=f"key_{i + 1:04d}") - query.add_range(r) - end_time = time.perf_counter() - - elapsed = end_time - start_time - # Under O(N^2) on test runners, 1k iterations takes seconds. - # Under O(1) with tuple hashing, it takes <0.10s. - # Double the value to avoid flaky tests - assert elapsed < 0.20, ( - f"O(N^2) Performance regression detected! Inserting {ranges_count} ranges took {elapsed:.2f}s" - ) + + # Intercept the `row_ranges` class property with a mock to track accesses + with mock.patch.object( + type(query), "row_ranges", new_callable=mock.PropertyMock + ) as row_ranges: + for i in range(100): + # Add unique sequential ranges. Utilizing our fast O(1) set signature lookup + # should never touch the slow wrapper-materialization properties. + query.add_range(RowRange(f"k{i:03d}", f"k{i + 1:03d}")) + + # ASSERT: The slow, O(N) wrapper-building property was never accessed + row_ranges.assert_not_called() + + # Verify that all ranges were successfully committed to the underlying protobuf + assert len(query._row_set.row_ranges) == 100 def test_add_range_deduplication_behavior(self): """Assert that identical ranges are deduplicated, but distinct ranges are kept.""" From b864707aa5f0885f0b0b59cb4f5dc4642bc9a164 Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Thu, 15 Oct 2026 21:20:11 +0000 Subject: [PATCH 6/6] move import --- .../tests/unit/data/test_read_rows_query.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py b/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py index 66eaa26663ad..703f9118646d 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import time +from unittest import mock import pytest @@ -602,7 +602,6 @@ def test_add_range_does_not_materialize_row_ranges(self): that it is never accessed during batch additions, proving we bypass the quadratic evaluation path entirely. """ - from unittest import mock from google.cloud.bigtable.data.read_rows_query import RowRange query = self._make_one()