fix(bigtable): resolve O(n^2) bottleneck in ReadRowsQuery.add_range

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:

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