Skip to content
Original file line number Diff line number Diff line change
Expand Up @@ -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:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

__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

return hash(self._signature())

@property
def start_key(self) -> bytes | None:
"""
Expand Down Expand Up @@ -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

Expand All @@ -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

Expand Down Expand Up @@ -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
Comment on lines +268 to +269

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

)

self.limit = limit or None
self.filter = row_filter

Expand Down Expand Up @@ -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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down Expand Up @@ -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):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

"""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
Loading