This is an automated email from the ASF dual-hosted git repository.
shunping pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new eb9f10482b0 Fix Bigtable DirectRow mutation protobuf retrieval for
google-cloud-bigtable >= 2.44.0 (#40029)
eb9f10482b0 is described below
commit eb9f10482b0c7a462a4edcfeba5a809ae5f9a099
Author: Shunping Huang <[email protected]>
AuthorDate: Sun Sep 6 07:36:03 2026 -0400
Fix Bigtable DirectRow mutation protobuf retrieval for
google-cloud-bigtable >= 2.44.0 (#40029)
---
sdks/python/apache_beam/io/gcp/bigtableio.py | 8 +++++++-
sdks/python/apache_beam/io/gcp/bigtableio_test.py | 17 +++++++++++++----
2 files changed, 20 insertions(+), 5 deletions(-)
diff --git a/sdks/python/apache_beam/io/gcp/bigtableio.py
b/sdks/python/apache_beam/io/gcp/bigtableio.py
index 442ba5d68c8..cd78deb7466 100644
--- a/sdks/python/apache_beam/io/gcp/bigtableio.py
+++ b/sdks/python/apache_beam/io/gcp/bigtableio.py
@@ -283,7 +283,13 @@ class WriteToBigTable(beam.PTransform):
def process(self, direct_row):
args = {"key": direct_row.row_key, "mutations": []}
# start accumulating mutations in a list
- for mutation in direct_row._get_mutations():
+ # In google-cloud-bigtable >= 2.44.0, _get_mutations() returns Python
+ # dataclass objects (RowMutationEntry) instead of protobuf messages.
+ # Use _get_mutation_pbs() to retrieve Mutation protobuf objects.
+ mutations = (
+ direct_row._get_mutation_pbs() if hasattr(
+ direct_row, '_get_mutation_pbs') else
direct_row._get_mutations())
+ for mutation in mutations:
if mutation.__contains__("set_cell"):
mutation_dict = {
"type": b'SetCell',
diff --git a/sdks/python/apache_beam/io/gcp/bigtableio_test.py
b/sdks/python/apache_beam/io/gcp/bigtableio_test.py
index 7c371c9b383..08c33017f9c 100644
--- a/sdks/python/apache_beam/io/gcp/bigtableio_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigtableio_test.py
@@ -119,6 +119,15 @@ class TestBeamRowToPartialRowData(unittest.TestCase):
class TestBigtableDirectRowToBeamRow(unittest.TestCase):
doFn = bigtableio.WriteToBigTable._DirectRowMutationsToBeamRow()
+ @staticmethod
+ def _get_mutation_pbs(direct_row):
+ # In google-cloud-bigtable >= 2.44.0, _get_mutations() returns Python
+ # dataclass objects instead of protobuf messages; use _get_mutation_pbs()
+ # to retrieve Mutation protobuf messages.
+ if hasattr(direct_row, '_get_mutation_pbs'):
+ return direct_row._get_mutation_pbs()
+ return direct_row._get_mutations()
+
def test_set_cell(self):
# create some set cell mutations
direct_row: DirectRow = DirectRow('key-1')
@@ -144,7 +153,7 @@ class TestBigtableDirectRowToBeamRow(unittest.TestCase):
# sort both lists of mutations for convenience
beam_row_mutations = sorted(beam_row.mutations, key=lambda m: m['value'])
bt_row_mutations = sorted(
- direct_row._get_mutations(), key=lambda m: m.set_cell.value)
+ self._get_mutation_pbs(direct_row), key=lambda m: m.set_cell.value)
self.assertEqual(beam_row.key, direct_row.row_key)
self.assertEqual(len(beam_row_mutations), len(bt_row_mutations))
@@ -186,7 +195,7 @@ class TestBigtableDirectRowToBeamRow(unittest.TestCase):
beam_row_mutations = sorted(
beam_row.mutations, key=lambda m: m['column_qualifier'])
bt_row_mutations = sorted(
- direct_row._get_mutations(),
+ self._get_mutation_pbs(direct_row),
key=lambda m: m.delete_from_column.column_qualifier)
self.assertEqual(beam_row.key, direct_row.row_key)
self.assertEqual(len(beam_row_mutations), len(bt_row_mutations))
@@ -232,8 +241,8 @@ class TestBigtableDirectRowToBeamRow(unittest.TestCase):
beam_row_mutations = sorted(
beam_row.mutations, key=lambda m: m['family_name'])
bt_row_mutations = sorted(
- direct_row._get_mutations(),
- key=lambda m: m.delete_from_column.family_name)
+ self._get_mutation_pbs(direct_row),
+ key=lambda m: m.delete_from_family.family_name)
self.assertEqual(beam_row.key, direct_row.row_key)
self.assertEqual(len(beam_row_mutations), len(bt_row_mutations))
