diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py index 3e6c2b533b1e..d6fcfc73520b 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py @@ -88,7 +88,7 @@ class _ReadRowsOperationAsync: def __init__( self, - query: ReadRowsQuery, + query: ReadRowsQuery | dict[str, object], target: TargetType, operation_timeout: float, attempt_timeout: float, diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py index ce29226cac53..f8c2d28b2903 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -82,7 +82,7 @@ class _ReadRowsOperation: def __init__( self, - query: ReadRowsQuery, + query: ReadRowsQuery | dict[str, object], target: TargetType, operation_timeout: float, attempt_timeout: float, 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..938303ea5a4c 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,19 @@ def __init__( init_dict["end_key_open"] = end_key self._pb = RowRangePB(**init_dict) + def _key( + self, + ) -> tuple[bytes | None, bytes | None, bool | None, bool | None]: + return ( + self.start_key, + self.end_key, + self.start_is_inclusive, + self.end_is_inclusive, + ) + + def __hash__(self) -> int: + return hash(self._key()) + @property def start_key(self) -> bytes | None: """ @@ -229,7 +242,12 @@ class ReadRowsQuery: row_filter: a RowFilter to apply to the query """ - slots = ("_limit", "_filter", "_row_set") + __slots__ = ( + "_limit", + "_filter", + "_row_keys", + "_row_ranges", + ) def __init__( self, @@ -246,10 +264,13 @@ def __init__( row_ranges = [row_ranges] if not isinstance(row_keys, list): row_keys = [row_keys] - row_keys = [key.encode() if isinstance(key, str) else key for key in row_keys] - self._row_set = RowSetPB( - row_keys=row_keys, row_ranges=[r._pb for r in row_ranges] - ) + # Use dicts as ordered sets to preserve insertion order with O(1) deduplication + self._row_keys: dict[bytes, None] = {} + self._row_ranges: dict[RowRange, None] = {} + for k in row_keys: + self.add_key(k) + for r in row_ranges: + self.add_range(r) self.limit = limit or None self.filter = row_filter @@ -261,7 +282,7 @@ def row_keys(self) -> list[bytes]: Returns: list[bytes]: the row keys in this query """ - return list(self._row_set.row_keys) + return list(self._row_keys) @property def row_ranges(self) -> list[RowRange]: @@ -271,7 +292,7 @@ def row_ranges(self) -> list[RowRange]: Returns: list[RowRange]: the row ranges in this query """ - return [RowRange._from_pb(r) for r in self._row_set.row_ranges] + return list(self._row_ranges) @property def limit(self) -> int | None: @@ -336,8 +357,7 @@ def add_key(self, row_key: str | bytes): row_key = row_key.encode() elif not isinstance(row_key, bytes): raise ValueError("row_key must be string or bytes") - if row_key not in self._row_set.row_keys: - self._row_set.row_keys.append(row_key) + self._row_keys[row_key] = None def add_range( self, @@ -349,8 +369,7 @@ def add_range( Args: row_range: a range of row keys to add to this query """ - if row_range not in self.row_ranges: - self._row_set.row_ranges.append(row_range._pb) + self._row_ranges[row_range] = None def shard(self, shard_keys: RowKeySamples) -> ShardedQuery: """ @@ -491,10 +510,19 @@ def _to_pb(self, table) -> ReadRowsRequestPB: app_profile_id=table.app_profile_id, filter=self.filter._to_pb() if self.filter else None, rows_limit=self.limit or 0, - rows=self._row_set, + rows=RowSetPB( + row_keys=self.row_keys, + row_ranges=[r._pb for r in self._row_ranges], + ), **table._request_path, ) + def __copy__(self) -> ReadRowsQuery: + new_query = ReadRowsQuery(limit=self.limit, row_filter=self.filter) + new_query._row_keys = self._row_keys.copy() + new_query._row_ranges = self._row_ranges.copy() + return new_query + def __eq__(self, other): """ RowRanges are equal if they have the same row keys, row ranges, @@ -509,24 +537,18 @@ def __eq__(self, other): if not isinstance(other, ReadRowsQuery): return False # empty queries are equal - if len(self.row_keys) == 0 and len(other.row_keys) == 0: - this_range_empty = len(self.row_ranges) == 0 or all( - [bool(r) is False for r in self.row_ranges] + if len(self._row_keys) == 0 and len(other._row_keys) == 0: + this_range_empty = len(self._row_ranges) == 0 or all( + [bool(r) is False for r in self._row_ranges] ) - other_range_empty = len(other.row_ranges) == 0 or all( - [bool(r) is False for r in other.row_ranges] + other_range_empty = len(other._row_ranges) == 0 or all( + [bool(r) is False for r in other._row_ranges] ) if this_range_empty and other_range_empty: return self.filter == other.filter and self.limit == other.limit - # otherwise, sets should have same sizes - if len(self.row_keys) != len(other.row_keys): - return False - if len(self.row_ranges) != len(other.row_ranges): - return False - ranges_match = all([row in other.row_ranges for row in self.row_ranges]) return ( - self.row_keys == other.row_keys - and ranges_match + self._row_keys.keys() == other._row_keys.keys() + and self._row_ranges.keys() == other._row_ranges.keys() and self.filter == other.filter and self.limit == other.limit ) 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..6e041e3ed7b8 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py @@ -157,17 +157,6 @@ class RowRange(_MappableAttributesMixin, BaseRowRange): "end_inclusive": "end_is_inclusive", } - def _key(self): - return ( - self.start_key, - self.end_key, - self.start_is_inclusive, - 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..cfef2da4b793 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,148 @@ 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): + """ + Verify that add_range deduplicates in O(1) time without accessing the + row_ranges property list. + """ + from google.cloud.bigtable.data.read_rows_query import RowRange + + query = self._make_one() + + with mock.patch.object( + type(query), "row_ranges", new_callable=mock.PropertyMock + ) as row_ranges: + for i in range(100): + query.add_range(RowRange(f"k{i:03d}", f"k{i + 1:03d}")) + row_ranges.assert_not_called() + assert len(query.row_ranges) == 100 + + def test_add_key_does_not_materialize_row_keys(self): + """ + Verify that ReadRowsQuery deduplicates row_keys in __init__ and add_key + in O(1) time without materializing the row_keys property list. + """ + query = self._make_one(row_keys=["k0", b"k0"]) + assert query.row_keys == [b"k0"] + + with mock.patch.object( + type(query), "row_keys", new_callable=mock.PropertyMock + ) as row_keys_prop: + query.add_key(b"k0") + query.add_key("k0") + row_keys_prop.assert_not_called() + + assert query.row_keys == [b"k0"] + + def test_query_equality_does_not_materialize_row_keys_or_ranges(self): + """ + Verify that ReadRowsQuery.__eq__ compares keys and ranges in order-independent + O(N) time using internal sets without materializing row_keys or row_ranges lists. + """ + from google.cloud.bigtable.data.read_rows_query import RowRange + + q1 = self._make_one( + row_keys=[f"k{i}" for i in range(50)], + row_ranges=[RowRange(f"a{i}", f"b{i}") for i in range(50)], + ) + q2 = self._make_one( + row_keys=[f"k{i}" for i in reversed(range(50))], + row_ranges=[RowRange(f"a{i}", f"b{i}") for i in reversed(range(50))], + ) + with ( + mock.patch.object( + type(q1), "row_keys", new_callable=mock.PropertyMock + ) as mock_keys, + mock.patch.object( + type(q1), "row_ranges", new_callable=mock.PropertyMock + ) as mock_ranges, + ): + assert q1 == q2 + mock_keys.assert_not_called() + mock_ranges.assert_not_called() + + 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 + + r1 = RowRange(start_key="a", end_key="b") + r2 = RowRange(start_key="a", end_key="b") + + r3 = RowRange(start_key="c", end_key="d") + + query = self._make_one(row_ranges=[r1]) + query.add_range(r2) + query.add_range(r3) + query.add_range(RowRange(start_key="c", end_key="d")) + + 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 + + base_range = RowRange(b"a", b"b") + legacy_range = LegacyRowRange(b"a", b"b") + assert base_range == legacy_range + assert LegacyRowRange._key is RowRange._key + assert LegacyRowRange.__hash__ is RowRange.__hash__ + assert hash(base_range) == hash(legacy_range) + query = self._make_one(row_ranges=[base_range]) + query.add_range(legacy_range) + assert len(query.row_ranges) == 1 + + def test_read_rows_query_enforces_slots(self): + """ + Verify that ReadRowsQuery enforces __slots__ so instances do not allocate + a dynamic __dict__ or permit undeclared attributes. + """ + from google.cloud.bigtable.data.read_rows_query import RowRange + + query = self._make_one( + row_keys=[b"k1"], + row_ranges=[RowRange(start_key=b"a", end_key=b"b")], + ) + assert not hasattr(query, "__dict__") + with pytest.raises(AttributeError): + query._undeclared_attribute = "invalid" + + def test_copy_query_preserves_deduplication_isolation(self): + """ + Verify that shallow-copying a ReadRowsQuery produces independent + deduplication tracking sets so mutations on the copy do not affect the + original query, and existing items remain deduplicated on the copy. + """ + import copy + + from google.cloud.bigtable.data.read_rows_query import RowRange + + r1 = RowRange(start_key="a", end_key="b") + r2 = RowRange(start_key="c", end_key="d") + + original = self._make_one(row_keys=[b"k1"], row_ranges=[r1]) + copied = copy.copy(original) + + copied.add_key(b"k1") + copied.add_range(r1) + assert copied.row_keys == [b"k1"] + assert copied.row_ranges == [r1] + + copied.add_key(b"k2") + copied.add_range(r2) + assert original.row_keys == [b"k1"] + assert original.row_ranges == [r1] + + original.add_key(b"k2") + original.add_range(r2) + + assert original.row_keys == [b"k1", b"k2"] + assert original.row_ranges == [r1, r2] diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py index 433444e0b8a9..af1ed4004928 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py @@ -1299,41 +1299,41 @@ def test__create_row_request_row_range_row_set_conflict(): def test__create_row_request_row_range_start_key(): + from google.cloud.bigtable.data.read_rows_query import RowRange from google.cloud.bigtable.table import _create_row_request - from google.cloud.bigtable_v2.types import RowRange start_key = b"begin_key" result = _create_row_request(start_key=start_key) - expected_result = ReadRowsQuery(row_ranges=RowRange(start_key_closed=start_key)) + expected_result = ReadRowsQuery(row_ranges=RowRange(start_key=start_key)) assert result == expected_result def test__create_row_request_row_range_end_key(): + from google.cloud.bigtable.data.read_rows_query import RowRange from google.cloud.bigtable.table import _create_row_request - from google.cloud.bigtable_v2.types import RowRange end_key = b"begin_key" result = _create_row_request(end_key=end_key) - expected_result = ReadRowsQuery(row_ranges=RowRange(end_key_open=end_key)) + expected_result = ReadRowsQuery(row_ranges=RowRange(end_key=end_key)) assert result == expected_result def test__create_row_request_row_range_both_keys(): + from google.cloud.bigtable.data.read_rows_query import RowRange from google.cloud.bigtable.table import _create_row_request - from google.cloud.bigtable_v2.types import RowRange start_key = b"begin_key" end_key = b"end_key" result = _create_row_request(start_key=start_key, end_key=end_key) expected_result = ReadRowsQuery( - row_ranges=RowRange(start_key_closed=start_key, end_key_open=end_key) + row_ranges=RowRange(start_key=start_key, end_key=end_key) ) assert result == expected_result def test__create_row_request_row_range_both_keys_inclusive(): + from google.cloud.bigtable.data.read_rows_query import RowRange from google.cloud.bigtable.table import _create_row_request - from google.cloud.bigtable_v2.types import RowRange start_key = b"begin_key" end_key = b"end_key" @@ -1341,7 +1341,7 @@ def test__create_row_request_row_range_both_keys_inclusive(): start_key=start_key, end_key=end_key, end_inclusive=True ) expected_result = ReadRowsQuery( - row_ranges=RowRange(start_key_closed=start_key, end_key_closed=end_key) + row_ranges=RowRange(start_key=start_key, end_key=end_key, end_is_inclusive=True) ) assert result == expected_result