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..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 @@ -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: """ @@ -140,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 @@ -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..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,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest import mock + import pytest TEST_ROWS = [ @@ -586,3 +588,58 @@ 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_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 google.cloud.bigtable.data.read_rows_query import RowRange + + query = self._make_one() + + # 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 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