fix(bigtable): resolve O(n^2) bottleneck in ReadRowsQuery.add_range by parthea · Pull Request #18706 · googleapis/google-cloud-python

parthea · 2026-10-15T19:51:04Z

This PR resolves a severe client-side O(n^2) CPU performance bottleneck inside ReadRowsQuery.add_range. This regression was introduced when deduplication checks were added to prevent duplicate row ranges in read requests.

Fixes b/573522284
Fixes #18708

gemini-code-assist

Code Review

This pull request optimizes the performance of ReadRowsQuery.add_range from O(N) to O(1) by making RowRange hashable and maintaining a _range_signatures set. It also fixes a typo from slots to __slots__ to ensure proper memory optimization, and adds corresponding unit tests. However, the proposed self-healing logic in add_range compares the lengths of the _range_signatures set and the _row_set.row_ranges list. If the query is initialized with duplicate ranges, these lengths will never match, causing the cache to rebuild on every call and introducing an O(N^2) bottleneck. It is recommended to remove this self-healing check entirely.

ohmayr

Thanks for picking this up. Using a set is the right fix. One blocking issue (inline): __hash__ and __eq__ disagree for the legacy row_set.RowRange, so mixed legacy/base ranges no longer dedup. The rest are small.

Separately, add_key and ReadRowsQuery.__eq__ have the same O(n²) pattern: 10k add_key calls take ~11.6s, and __eq__ on two 2k-range queries takes ~50s. These can go in a follow-up PR, but please don't close b/573522284 until they're tracked, since key-heavy workloads are still slow.

ohmayr · 2026-10-15T21:19:56Z

+            self.end_is_inclusive,
+        )
+
+    def __hash__(self) -> int:


__eq__ compares _pb, but __hash__ hashes self._key(). The legacy subclass google.cloud.bigtable.row_set.RowRange overrides _key() with a different field order, so legacy and base ranges compare equal but hash differently. That breaks set membership for mixed types, and dedup regresses vs main:

from google.cloud.bigtable.row_set import RowRange as LegacyRowRange b, l = RowRange(b"a", b"b"), LegacyRowRange(b"a", b"b") b == l # True hash(b) == hash(l) # False q = ReadRowsQuery(row_ranges=[b]) q.add_range(l) len(q.row_ranges) # 2 on this branch, 1 on main

Suggest hashing the same data __eq__ compares:

def __hash__(self) -> int: pb = self._pb return hash( (pb.start_key_closed, pb.start_key_open, pb.end_key_closed, pb.end_key_open) )

and dropping __hash__ from row_set.RowRange so it inherits this one (_key() then isn't needed here). I tried this locally: mixed legacy/base dedup works, the bigtable unit tests pass, and 10k add_range calls still take ~0.23s. Could you add a test along these lines?

def test_row_range_hash_matches_legacy(self): from google.cloud.bigtable.data.read_rows_query import RowRange from google.cloud.bigtable.row_set import RowRange as LegacyRowRange b, l = RowRange(b"a", b"b"), LegacyRowRange(b"a", b"b") assert b == l assert hash(b) == hash(l) query = self._make_one(row_ranges=[b]) query.add_range(l) assert len(query.row_ranges) == 1

Done in ef09fca.

RowRange.__hash__ now hashes the 4 proto fields (via _signature()), row_set.RowRange.__hash__ is removed so it inherits the base implementation, and I added test_row_range_hash_matches_legacy.

ohmayr · 2026-10-15T21:19:56Z

+        self._range_signatures = set(
+            RowRange._from_pb(r) for r in self._row_set.row_ranges


row_ranges is already a list of RowRange here, so set(row_ranges) should be enough and avoids re-wrapping from the proto. Naming nit: this holds RowRange objects, not signatures. Maybe _row_range_set?

Heads-up: ReadRowsQuery declares slots = (...) (not __slots__), which is the only reason adding a new attribute works. If that typo gets fixed later, this attribute needs to be added there.

Addressed in f9b35b1 and 213fa5b: __init__ now calls self.add_range(r) directly, _range_signatures stores the immutable row_range._signature() tuple, and slots is now __slots__ with the new attributes included.

ohmayr · 2026-10-15T21:19:56Z

+        # 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):


Thanks for adding this. One gap: the new set seeding in __init__ isn't covered. Could we add a case where a range is passed to the constructor and an equal one is then added?

query = self._make_one(row_ranges=[RowRange("a", "b")]) query.add_range(RowRange("a", "b")) assert len(query.row_ranges) == 1

Done in bff7b0b

ohmayr · 2026-10-15T21:19:57Z

            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


nit: I'd drop this comment. The code already says what it does.

Fixed in f9b35b1

ohmayr · 2026-10-15T21:19:57Z

    def _from_pb(cls, data: RowRangePB) -> RowRange:
        """
-        Creates a RowRange from a protobuf
+        Creates a RowRange from a protobuf without allocating temporary unused objects.


nit: cls.__new__(cls) is a good change, but the docstring now describes implementation. I'd keep the original wording and add a # comment on why __init__ is skipped if needed.

Done in 183df93

…gacy classes

…and deduplicate in __init__

…nforce slots

parthea · 2026-10-19T19:35:47Z

Thanks for picking this up. Using a set is the right fix. One blocking issue (inline): __hash__ and __eq__ disagree for the legacy row_set.RowRange, so mixed legacy/base ranges no longer dedup. The rest are small.

Separately, add_key and ReadRowsQuery.__eq__ have the same O(n²) pattern: 10k add_key calls take ~11.6s, and __eq__ on two 2k-range queries takes ~50s. These can go in a follow-up PR, but please don't close b/573522284 until they're tracked, since key-heavy workloads are still slow.

I've addressed both points directly in this PR (we can fully resolve b/573522284):

Fixed in ef09fca (and updated _range_signatures to store the tuple directly in f9b35b1).

I also fixed the performance issue in 213fa5b

…ignature

… tests

…gePB)

parthea · 2026-10-19T20:50:11Z

/gemini review

gemini-code-assist

Code Review

This pull request optimizes the performance of ReadRowsQuery and RowRange by introducing __slots__ to reduce memory overhead, implementing O(1) deduplication for row keys and ranges using internal sets (_key_set and _range_signatures), and optimizing equality checks to avoid materializing wrapper objects. It also updates type annotations and adds comprehensive unit tests to verify these performance improvements and deduplication behaviors. Feedback suggests optimizing the __copy__ method of ReadRowsQuery by using the native protobuf CopyFrom method via the public API type(value).pb(value) instead of copy.deepcopy on the private _pb attribute to further improve performance and avoid accessing private fields.

parthea · 2026-10-19T20:59:11Z

/gemini review

gemini-code-assist

Code Review

This pull request optimizes ReadRowsQuery and RowRange by introducing O(1) set-based deduplication for row keys and ranges, enforcing __slots__ to reduce memory allocation, and optimizing object copying and instantiation. The review feedback suggests simplifying the nested conditional logic in _signature_from_pb by using the protobuf WhichOneof method to improve readability and idiomatic handling of the oneof fields.

gemini-code-assist · 2026-10-19T21:02:42Z

+        # 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


The current nested conditional logic with HasField checks is complex and prone to errors when handling edge cases like empty bytes (b""). Since start_key and end_key are defined as oneof fields in the protobuf definition, we can use the standard WhichOneof method on the underlying protobuf message. This simplifies the logic significantly, improves readability, and is more idiomatic.

start_oneof = raw_pb.WhichOneof("start_key") if start_oneof == "start_key_closed": start_key = raw_pb.start_key_closed start_is_inclusive = True elif start_oneof == "start_key_open": start_key = raw_pb.start_key_open start_is_inclusive = False else: start_key = None start_is_inclusive = True end_oneof = raw_pb.WhichOneof("end_key") if end_oneof == "end_key_open": end_key = raw_pb.end_key_open end_is_inclusive = False elif end_oneof == "end_key_closed": end_key = raw_pb.end_key_closed end_is_inclusive = True else: end_key = None end_is_inclusive = False

Keeping the current structure here intentionally as a hot-path optimization.

See analysis below:

In the common case (non-empty start_key_closed and end_key_open), checking truthiness directly on raw_pb.start_key_closed and raw_pb.end_key_open avoids extra WhichOneof method call and string comparison overhead on every add_range call, only falling back to HasField when a bound is empty (b"") or unset (None).

IIUC, we shouldn't need the HasField fallback case at all. Bigtable treats falsy values as unbounded keys, and the handwritten classes collapse them into None

daniel-sanche · 2026-10-19T20:55:28Z

+        # 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.


nit: Is this comment really useful? It seems like it would feel out of place after the PR is merged

daniel-sanche · 2026-10-19T21:01:13Z

+    @staticmethod
+    def _signature_from_pb(
+        pb: RowRangePB,
+    ) -> tuple[bytes | None, bool, bytes | None, bool]:


nit: if we do keep this, we should probably add a type alias for the signature

daniel-sanche · 2026-10-19T21:03:43Z

    def add_range(
        self,
-        row_range: RowRange,
+        row_range: RowRange | RowRangePB,


Is this necessary? Bigtable wanted to hide protobufs from the public interface

daniel-sanche · 2026-10-19T21:21:48Z

+        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)


_signature_from_pb seems very complex. Isn't this doing the same job as the legacy client's _key()? Can't we just move _key() and __hash__() into the base class here, and avoid the complex parsing logic?

(Also, are you sure we need both _signature and _signature_from_pb? I only see _signature_from_pb being used to handle PB inputs, but I don't think we'd actually want to expose that publicly)

daniel-sanche · 2026-10-19T21:39:35Z

-        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()


It seems like we're duplicating some state now, if we are tracking both row_set, and keys/ranges separately. Lots of room for those to get out of sync

Could we remove row_set, turn range_signatures into a set of full Range instances, and just re-construct the row_set inside to_pb? That seems like it could simplify the data model. Or am I misunderstanding?

daniel-sanche · 2026-10-19T22:03:36Z

+        # 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


IIUC, we shouldn't need the HasField fallback case at all. Bigtable treats falsy values as unbounded keys, and the handwritten classes collapse them into None

gemini-code-assist Bot reviewed Oct 15, 2026

View reviewed changes

Comment thread packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py Outdated

parthea force-pushed the fix-bigtable-20261015 branch from c486b34 to f6a2fa6 Compare October 15, 2026 19:55

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

f8d3b11

parthea force-pushed the fix-bigtable-20261015 branch from f6a2fa6 to f8d3b11 Compare October 15, 2026 20:03

address offline feedback

679ebed

parthea marked this pull request as ready for review October 15, 2026 20:15

parthea requested a review from a team as a code owner October 15, 2026 20:15

update comment

285b2c1

sahananshylaja-glitch mentioned this pull request Oct 15, 2026

perf(bigtable): resolve O(n^2) performance bottleneck in ReadRowsQuery.add_range #18705

Closed

2 of 4 tasks

fix build; add test

01e759a

parthea mentioned this pull request Oct 15, 2026

bug(bigtable): Performance regression in ReadRowsQuery.add_range #18708

Open

ohmayr reviewed Oct 15, 2026

View reviewed changes

Comment thread packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py Outdated

address review feedback

d08cc34

ohmayr reviewed Oct 15, 2026

View reviewed changes

move import

b864707

parthea self-assigned this Oct 16, 2026

parthea added 8 commits October 17, 2026 13:38

Merge remote-tracking branch 'origin/main' into fix-bigtable-20261015

3f54cac

fix(bigtable): align RowRange.__hash__ with __eq__ across base and le…

ef09fca

…gacy classes

address review feedback

183df93

address offline feedback

39ddcf0

fix(bigtable): store immutable tuple signatures in _range_signatures …

f9b35b1

…and deduplicate in __init__

perf(bigtable): resolve O(n^2) in ReadRowsQuery.add_key and __eq__, e…

213fa5b

…nforce slots

fix build

bb51980

mypy

d2b8192

parthea added 5 commits October 19, 2026 19:45

address review feedback

bff7b0b

lint

09110b3

Merge remote-tracking branch 'origin/main' into fix-bigtable-20261015

26ac51f

fix(bigtable): distinguish unset and empty-byte bounds in RowRange._s…

73f1e1c

…ignature

test(bigtable): tighten mock and mutation assertions in ReadRowsQuery…

6d94d4f

… tests

parthea added 2 commits October 19, 2026 20:41

perf(bigtable): enforce __slots__ on legacy row_set.RowRange

8e32f9c

perf(bigtable): avoid RowRange wrapper allocation in add_range(RowRan…

3491e36

…gePB)

run nox -s generate_sync

d88fecb

gemini-code-assist Bot reviewed Oct 19, 2026

View reviewed changes

Comment thread packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py

address review feedback

0f285aa

gemini-code-assist Bot reviewed Oct 19, 2026

View reviewed changes

lint

64125aa

daniel-sanche requested changes Oct 19, 2026

View reviewed changes

		self._range_signatures = set(
		RowRange._from_pb(r) for r in self._row_set.row_ranges

Conversation

parthea commented Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

ohmayr left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

parthea commented Oct 19, 2026

Uh oh!

parthea commented Oct 19, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

parthea commented Oct 19, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Oct 19, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

parthea commented Oct 15, 2026 •

edited

Loading