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..f9f82b38d733 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,18 @@ def __init__( init_dict["end_key_open"] = end_key self._pb = RowRangePB(**init_dict) + def _signature(self) -> tuple[bytes, bytes, bytes, bytes]: + pb = self._pb + return ( + pb.start_key_closed, + pb.start_key_open, + pb.end_key_closed, + pb.end_key_open, + ) + + def __hash__(self) -> int: + return hash(self._signature()) + @property def start_key(self) -> bytes | None: """ @@ -147,7 +159,8 @@ def _from_pb(cls, data: RowRangePB) -> RowRange: Returns: RowRange: The converted RowRange """ - instance = cls() + # Skip __init__ to avoid allocating an unused RowRangePB + instance = cls.__new__(cls) instance._pb = data return instance @@ -164,7 +177,8 @@ def _from_dict(cls, data: dict[str, bytes | str]) -> RowRange: formatted_data = { k: v.encode() if isinstance(v, str) else v for k, v in data.items() } - instance = cls() + # Skip __init__ to avoid allocating an unused RowRangePB + instance = cls.__new__(cls) instance._pb = RowRangePB(**formatted_data) return instance @@ -250,6 +264,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 +368,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/google/cloud/bigtable/row_set.py b/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py index 4e4bacfa7cda..0732b3e2d312 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py @@ -165,9 +165,6 @@ def _key(self): self.end_is_inclusive, ) - def __hash__(self): - return hash(self._key()) - def get_range_kwargs(self): """Convert row range object to dict which can be passed to google.bigtable.v2.RowRange add method. 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..2a8f51dc49ce 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,82 @@ 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 + + def test_row_range_hash_matches_legacy(self): + """ + Verify that base RowRange and legacy row_set.RowRange instances representing + the same range compare equal, produce identical hashes, and deduplicate + properly when combined in a ReadRowsQuery. + """ + from google.cloud.bigtable.data.read_rows_query import RowRange + from google.cloud.bigtable.row_set import RowRange as LegacyRowRange + + # Instantiate equivalent ranges using both the data API and legacy classes. + base_range = RowRange(b"a", b"b") + legacy_range = LegacyRowRange(b"a", b"b") + + # Equal RowRange objects must produce identical hashes so set membership + # operates consistently across base and legacy types. + assert base_range == legacy_range + assert hash(base_range) == hash(legacy_range) + + # Adding an equivalent legacy range to a query initialized with a base range + # should be recognized as a duplicate and ignored. + query = self._make_one(row_ranges=[base_range]) + query.add_range(legacy_range) + assert len(query.row_ranges) == 1