Skip to content

Commit e688531

Browse files
mutianfaxyjogemini-code-assist[bot]
authored
bigtable: add ValueBitmaskFilter for data client (#17567)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/google-cloud-python/issues) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #<issue_number_goes_here> 🦕 --------- Co-authored-by: Akshay Joshi <akshay@anthropic.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 parent 8cb77d9 commit e688531

4 files changed

Lines changed: 115 additions & 0 deletions

File tree

packages/google-cloud-bigtable/google/cloud/bigtable/data/row_filters.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -484,6 +484,34 @@ def _to_dict(self) -> dict[str, bytes]:
484484
return {"value_regex_filter": self.regex}
485485

486486

487+
class ValueBitmaskFilter(RowFilter):
488+
"""Row filter for a value bitmask.
489+
490+
Matches only cells with values that satisfy the condition
491+
``(value & mask) == mask``. The mask length must exactly match the value
492+
length, otherwise the cell is not considered a match.
493+
494+
:type mask: bytes or str
495+
:param mask: A bitmask to match against cell values. String values
496+
will be encoded as ASCII.
497+
"""
498+
499+
def __init__(self, mask: bytes | str):
500+
self.mask: bytes = _to_bytes(mask)
501+
502+
def __eq__(self, other):
503+
if not isinstance(other, ValueBitmaskFilter):
504+
return NotImplemented
505+
return other.mask == self.mask
506+
507+
def _to_dict(self) -> dict[str, Any]:
508+
"""Converts the row filter to a dict representation."""
509+
return {"value_bitmask_filter": {"mask": self.mask}}
510+
511+
def __repr__(self) -> str:
512+
return f"{self.__class__.__name__}(mask={self.mask!r})"
513+
514+
487515
class LiteralValueFilter(ValueRegexFilter):
488516
"""Row filter for an exact value.
489517

packages/google-cloud-bigtable/tests/system/data/test_system_async.py

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1150,6 +1150,41 @@ async def test_literal_value_filter(
11501150
f"row {type(cell_value)}({cell_value}) not found with {type(filter_input)}({filter_input}) filter"
11511151
)
11521152

1153+
@pytest.mark.usefixtures("target")
1154+
@CrossSync.Retry(
1155+
predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
1156+
)
1157+
@pytest.mark.parametrize(
1158+
"cell_value,mask,expect_match",
1159+
[
1160+
(b"\x01\x02\x03", b"\x01\x02\x03", True),
1161+
(b"\x01\x02\x03", b"\x01\x00\x00", True),
1162+
(b"\x00\x02\x03", b"\x01\x00\x00", False),
1163+
],
1164+
)
1165+
@pytest.mark.skipif(
1166+
bool(os.environ.get(BIGTABLE_EMULATOR)),
1167+
reason="value_bitmask_filter not supported by emulator",
1168+
)
1169+
@CrossSync.pytest
1170+
async def test_value_bitmask_filter(
1171+
self, target, temp_rows, cell_value, mask, expect_match
1172+
):
1173+
"""
1174+
ValueBitmaskFilter matches cells where (value & mask) == mask.
1175+
Make sure inputs are properly interpreted by the server.
1176+
"""
1177+
from google.cloud.bigtable.data import ReadRowsQuery
1178+
from google.cloud.bigtable.data.row_filters import ValueBitmaskFilter
1179+
1180+
f = ValueBitmaskFilter(mask)
1181+
await temp_rows.add_row(b"row_key_1", value=cell_value)
1182+
query = ReadRowsQuery(row_keys=[b"row_key_1"], row_filter=f)
1183+
row_list = await target.read_rows(query)
1184+
assert len(row_list) == bool(expect_match), (
1185+
f"row {cell_value!r} not matched as {expect_match} with {mask!r} bitmask filter"
1186+
)
1187+
11531188
@pytest.mark.skipif(
11541189
bool(os.environ.get(BIGTABLE_EMULATOR)),
11551190
reason="emulator doesn't support SQL",

packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -928,6 +928,38 @@ def test_literal_value_filter(
928928
f"row {type(cell_value)}({cell_value}) not found with {type(filter_input)}({filter_input}) filter"
929929
)
930930

931+
@pytest.mark.usefixtures("target")
932+
@CrossSync._Sync_Impl.Retry(
933+
predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
934+
)
935+
@pytest.mark.parametrize(
936+
"cell_value,mask,expect_match",
937+
[
938+
(b"\x01\x02\x03", b"\x01\x02\x03", True),
939+
(b"\x01\x02\x03", b"\x01\x00\x00", True),
940+
(b"\x00\x02\x03", b"\x01\x00\x00", False),
941+
],
942+
)
943+
@pytest.mark.skipif(
944+
bool(os.environ.get(BIGTABLE_EMULATOR)),
945+
reason="value_bitmask_filter not supported by emulator",
946+
)
947+
def test_value_bitmask_filter(
948+
self, target, temp_rows, cell_value, mask, expect_match
949+
):
950+
"""ValueBitmaskFilter matches cells where (value & mask) == mask.
951+
Make sure inputs are properly interpreted by the server."""
952+
from google.cloud.bigtable.data import ReadRowsQuery
953+
from google.cloud.bigtable.data.row_filters import ValueBitmaskFilter
954+
955+
f = ValueBitmaskFilter(mask)
956+
temp_rows.add_row(b"row_key_1", value=cell_value)
957+
query = ReadRowsQuery(row_keys=[b"row_key_1"], row_filter=f)
958+
row_list = target.read_rows(query)
959+
assert len(row_list) == bool(expect_match), (
960+
f"row {cell_value!r} not matched as {expect_match} with {mask!r} bitmask filter"
961+
)
962+
931963
@pytest.mark.skipif(
932964
bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't support SQL"
933965
)

packages/google-cloud-bigtable/tests/unit/data/test_row_filters.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1987,6 +1987,26 @@ def test_literal_value__write_literal_regex(input_arg, expected_bytes):
19871987
assert filter_.regex == expected_bytes
19881988

19891989

1990+
class TestValueBitmaskFilter:
1991+
@staticmethod
1992+
def _target_class():
1993+
from google.cloud.bigtable.data.row_filters import ValueBitmaskFilter
1994+
1995+
return ValueBitmaskFilter
1996+
1997+
def test_to_dict(self):
1998+
mask = b"\xaa" * 8
1999+
row_filter = self._target_class()(mask)
2000+
expected = {"value_bitmask_filter": {"mask": mask}}
2001+
assert row_filter._to_dict() == expected
2002+
2003+
def test_to_pb(self):
2004+
mask = b"\xaa" * 8
2005+
row_filter = self._target_class()(mask)
2006+
pb = row_filter._to_pb()
2007+
assert pb.value_bitmask_filter.mask == mask
2008+
2009+
19902010
def _ColumnRangePB(*args, **kw):
19912011
from google.cloud.bigtable_v2.types import data as data_v2_pb2
19922012

0 commit comments

Comments
 (0)