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..4180d4a960d5 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, @@ -99,6 +99,10 @@ def __init__( attempt_timeout, operation_timeout ) self.operation_timeout = operation_timeout + # ReadRowsQuery defines __slots__, which gives it a C-level memory layout + # disjoint from dict. Consequently, `query` must be explicitly typed as + # `ReadRowsQuery | dict[str, object]` in the signature above; otherwise, + # mypy flags this runtime `isinstance(query, dict)` branch as unreachable. if isinstance(query, dict): self.request = ReadRowsRequestPB( **query, 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..e8731e2147de 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,59 @@ def __init__( init_dict["end_key_open"] = end_key self._pb = RowRangePB(**init_dict) + @staticmethod + def _signature_from_pb( + pb: RowRangePB, + ) -> tuple[bytes | None, bool, bytes | None, bool]: + raw_pb = pb._pb + + # Fast path: non-empty keys bypass HasField and proto-plus overhead + start_closed = raw_pb.start_key_closed + if start_closed: + start_key: bytes | None = start_closed + start_is_inclusive = True + else: + start_open = raw_pb.start_key_open + if start_open: + start_key = start_open + start_is_inclusive = False + elif raw_pb.HasField("start_key_closed"): + start_key = b"" + start_is_inclusive = True + elif raw_pb.HasField("start_key_open"): + start_key = b"" + start_is_inclusive = False + else: + start_key = None + start_is_inclusive = True + + end_open = raw_pb.end_key_open + if end_open: + end_key: bytes | None = end_open + end_is_inclusive = False + else: + end_closed = raw_pb.end_key_closed + if end_closed: + end_key = end_closed + end_is_inclusive = True + elif raw_pb.HasField("end_key_open"): + end_key = b"" + end_is_inclusive = False + elif raw_pb.HasField("end_key_closed"): + end_key = b"" + end_is_inclusive = True + else: + end_key = None + end_is_inclusive = False + + return (start_key, start_is_inclusive, end_key, end_is_inclusive) + + def _signature(self) -> tuple[bytes | None, bool, bytes | None, bool]: + return self._signature_from_pb(self._pb) + + def __hash__(self) -> int: + return hash(self._signature()) + @property def start_key(self) -> bytes | None: """ @@ -147,7 +200,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 +218,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 @@ -229,7 +284,13 @@ class ReadRowsQuery: row_filter: a RowFilter to apply to the query """ - slots = ("_limit", "_filter", "_row_set") + __slots__ = ( + "_limit", + "_filter", + "_row_set", + "_key_set", + "_range_signatures", + ) def __init__( self, @@ -246,10 +307,15 @@ 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] + self._row_set = RowSetPB() + self._key_set: set[bytes] = set() + self._range_signatures: set[tuple[bytes | None, bool, bytes | None, bool]] = ( + set() ) + 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 @@ -336,12 +402,13 @@ 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: + if row_key not in self._key_set: + self._key_set.add(row_key) self._row_set.row_keys.append(row_key) def add_range( self, - row_range: RowRange, + row_range: RowRange | RowRangePB, ): """ Add a range of row keys to this query. @@ -349,8 +416,15 @@ 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) + if isinstance(row_range, RowRangePB): + pb = row_range + sig = RowRange._signature_from_pb(pb) + else: + pb = row_range._pb + sig = row_range._signature() + if sig not in self._range_signatures: + self._range_signatures.add(sig) + self._row_set.row_ranges.append(pb) def shard(self, shard_keys: RowKeySamples) -> ShardedQuery: """ @@ -495,6 +569,15 @@ def _to_pb(self, table) -> ReadRowsRequestPB: **table._request_path, ) + def __copy__(self) -> ReadRowsQuery: + new_query = ReadRowsQuery(limit=self.limit, row_filter=self.filter) + type(self._row_set).pb(new_query._row_set).CopyFrom( + type(self._row_set).pb(self._row_set) + ) + new_query._key_set = self._key_set.copy() + new_query._range_signatures = self._range_signatures.copy() + return new_query + def __eq__(self, other): """ RowRanges are equal if they have the same row keys, row ranges, @@ -509,24 +592,24 @@ 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._key_set) == 0 and len(other._key_set) == 0: + empty_sig = (None, True, None, False) + this_range_empty = len(self._range_signatures) == 0 or ( + self._range_signatures == {empty_sig} ) - 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._range_signatures) == 0 or ( + other._range_signatures == {empty_sig} ) 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): + if len(self._key_set) != len(other._key_set): return False - if len(self.row_ranges) != len(other.row_ranges): + if len(self._range_signatures) != len(other._range_signatures): 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._key_set == other._key_set + and self._range_signatures == other._range_signatures and self.filter == other.filter and self.limit == other.limit ) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/helpers.py b/packages/google-cloud-bigtable/google/cloud/bigtable/helpers.py index 71af2f539f2c..933dd4614c4e 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/helpers.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/helpers.py @@ -43,6 +43,8 @@ class and define a map _attribute_map from legacy client attributes to data clie Attributes are remapped and redefined in __init__ as well as getattr/setattr. """ + __slots__ = () + def __init__(self, *args, **kwargs): new_kwargs = {self._attribute_map.get(k, k): v for (k, v) in kwargs.items()} super(_MappableAttributesMixin, self).__init__(*args, **new_kwargs) 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..d9a6bb08ab86 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py @@ -152,6 +152,8 @@ class RowRange(_MappableAttributesMixin, BaseRowRange): considered inclusive. The default is False (exclusive). """ + __slots__ = () + _attribute_map = { "start_inclusive": "start_is_inclusive", "end_inclusive": "end_is_inclusive", @@ -165,9 +167,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..da8cc3c3261e 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,337 @@ 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_key_deduplicates_via_set_without_scanning_row_set(self): + """ + Verify that ReadRowsQuery deduplicates row_keys in __init__ and add_key + using its internal O(1) key set without linearly scanning _row_set.row_keys. + """ + # Duplicate keys passed to __init__ (even in mixed str/bytes forms) + # must be deduplicated on construction. + query = self._make_one(row_keys=["k0", b"k0"]) + assert query.row_keys == [b"k0"] + + # When adding keys that are already present in the query, add_key must + # short-circuit via _key_set without accessing _row_set.row_keys. + mock_row_set = mock.MagicMock() + row_keys_prop = mock.PropertyMock() + type(mock_row_set).row_keys = row_keys_prop + with mock.patch.object(query, "_row_set", mock_row_set): + 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 + + # Create two equivalent queries with identical keys and ranges inserted + # in opposite orders. + 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))], + ) + + # Intercept the row_keys and row_ranges properties to verify that equality + # comparisons operate on internal sets without constructing wrapper lists. + 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 + + # 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") + + # Seed r1 via __init__, then add r2 (duplicate) and r3 (distinct) + query = self._make_one(row_ranges=[r1]) + query.add_range(r2) # Duplicate of constructor-seeded r1, should be ignored + query.add_range(r3) # Distinct, should be kept + query.add_range( + RowRange(start_key="c", end_key="d") + ) # Duplicate of r3, should be ignored + + # 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 + + def test_add_range_immune_to_caller_mutating_row_range_after_insert(self): + """ + Verify that mutating a RowRange's underlying protobuf after adding it to + a query does not corrupt the query's internal deduplication set. + """ + from google.cloud.bigtable.data.read_rows_query import RowRange + + query = self._make_one() + r1 = RowRange(start_key=b"a", end_key=b"b") + query.add_range(r1) + + # Mutate the caller's underlying RowRangePB in place after insertion. + r1._pb.start_key_closed = b"x" + r1._pb.end_key_open = b"y" + + # Adding a new RowRange with the original ("a", "b") bounds must still be + # detected as a duplicate of the previously added range. + query.add_range(RowRange(start_key=b"a", end_key=b"b")) + + assert len(query.row_ranges) == 1 + assert query.row_ranges[0] == RowRange(start_key=b"a", end_key=b"b") + + def test_row_range_from_dict_does_not_call_init(self): + """ + Verify that RowRange._from_dict constructs the wrapper instance directly + without invoking __init__ and allocating an unused intermediate protobuf. + """ + from google.cloud.bigtable.data.read_rows_query import RowRange + + # Spy on RowRange.__init__ to ensure dictionary deserialization bypasses + # default constructor initialization and dummy protobuf creation. + with mock.patch.object( + RowRange, "__init__", wraps=RowRange.__init__ + ) as mock_init: + row_range = RowRange._from_dict( + {"start_key_closed": b"a", "end_key_open": b"b"} + ) + mock_init.assert_not_called() + + # Verify the resulting RowRange exposes the expected boundaries and inclusivity. + assert row_range.start_key == b"a" + assert row_range.start_is_inclusive is True + assert row_range.end_key == b"b" + assert row_range.end_is_inclusive is False + + 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 + + # Initialize a query with both keys and ranges so all internal attributes + # are populated on the instance. + query = self._make_one( + row_keys=[b"k1"], + row_ranges=[RowRange(start_key=b"a", end_key=b"b")], + ) + + # Enforcing __slots__ prevents creation of a per-instance __dict__. + assert not hasattr(query, "__dict__") + + # Assigning an attribute not declared in __slots__ must raise AttributeError. + with pytest.raises(AttributeError): + query._undeclared_attribute = "invalid" + + def test_legacy_row_range_enforces_slots(self): + """ + Verify that legacy row_set.RowRange also defines __slots__ = () so it + inherits BaseRowRange's slotted layout without allocating a __dict__. + """ + from google.cloud.bigtable.row_set import RowRange as LegacyRowRange + + legacy_range = LegacyRowRange(start_key=b"a", end_key=b"b") + assert not hasattr(legacy_range, "__dict__") + + with pytest.raises(AttributeError): + legacy_range._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) + + # Re-adding existing items to the copy must be ignored as duplicates. + copied.add_key(b"k1") + copied.add_range(r1) + assert copied.row_keys == [b"k1"] + assert copied.row_ranges == [r1] + + # Mutate only the copied query by adding a new key and range. + copied.add_key(b"k2") + copied.add_range(r2) + + # Confirm the original query was not modified by operations on the copy. + assert original.row_keys == [b"k1"] + assert original.row_ranges == [r1] + + # Adding the same new key and range to the original query must succeed + # because its internal deduplication sets are independent from the copy. + original.add_key(b"k2") + original.add_range(r2) + + assert original.row_keys == [b"k1", b"k2"] + assert original.row_ranges == [r1, r2] + + def test_row_range_signature_distinguishes_empty_bytes_and_inclusivity(self): + """ + Verify that unset bounds (None), inclusive empty-byte bounds (b""), + and exclusive empty-byte bounds (b"") produce distinct signatures + and are not conflated during deduplication or query equality. + """ + from google.cloud.bigtable.data.read_rows_query import RowRange + + unbounded_start = RowRange(start_key=None, end_key=b"z") + closed_empty_start = RowRange._from_dict( + {"start_key_closed": b"", "end_key_open": b"z"} + ) + open_empty_start = RowRange._from_dict( + {"start_key_open": b"", "end_key_open": b"z"} + ) + # All three ranges are semantically distinct and must not collide + assert unbounded_start != closed_empty_start + assert closed_empty_start != open_empty_start + assert unbounded_start._signature() != closed_empty_start._signature() + assert closed_empty_start._signature() != open_empty_start._signature() + + query = self._make_one( + row_ranges=[unbounded_start, closed_empty_start, open_empty_start] + ) + assert len(query.row_ranges) == 3 + + # A range targeting the single row b"" must NOT compare equal to an empty query + single_empty_key_range = RowRange._from_dict( + {"start_key_closed": b"", "end_key_closed": b""} + ) + non_empty_query = self._make_one(row_ranges=[single_empty_key_range]) + empty_query = self._make_one() + assert non_empty_query != empty_query + + def test_add_range_pb_does_not_allocate_wrapper(self): + """ + Verify that passing a RowRangePB directly to add_range computes the + deduplication signature via _signature_from_pb without allocating a + temporary RowRange wrapper via _from_pb. + """ + from google.cloud.bigtable.data.read_rows_query import RowRange + from google.cloud.bigtable_v2.types import RowRange as RowRangePB + + query = self._make_one() + pb = RowRangePB(start_key_closed=b"a", end_key_open=b"b") + + with mock.patch.object( + RowRange, "_from_pb", wraps=RowRange._from_pb + ) as mock_from_pb: + query.add_range(pb) + query.add_range(pb) + mock_from_pb.assert_not_called() + + assert len(query.row_ranges) == 1 + assert query.row_ranges[0] == RowRange(start_key=b"a", end_key=b"b") + + def test_copy_query_uses_protobuf_copy_from_without_deepcopy(self): + """ + Verify that ReadRowsQuery.__copy__ clones _row_set using native + protobuf CopyFrom via type(_row_set).pb() rather than copy.deepcopy(). + """ + import copy + + from google.cloud.bigtable.data.read_rows_query import RowRange + from google.cloud.bigtable_v2.types import RowSet as RowSetPB + + r1 = RowRange(start_key=b"a", end_key=b"b") + original = self._make_one(row_keys=[b"k1"], row_ranges=[r1], limit=10) + + with ( + mock.patch("copy.deepcopy", wraps=copy.deepcopy) as mock_deepcopy, + mock.patch.object(RowSetPB, "pb", wraps=RowSetPB.pb) as mock_pb, + ): + copied = copy.copy(original) + mock_deepcopy.assert_not_called() + mock_pb.assert_any_call(original._row_set) + mock_pb.assert_any_call(copied._row_set) + + assert copied == original + assert copied._row_set is not original._row_set + assert copied.row_keys == [b"k1"] + assert copied.row_ranges == [r1]