From 4777abfd4f82337eadb822d6b4f9f4281e5d9987 Mon Sep 17 00:00:00 2001 From: aibrahiim Date: Mon, 7 Sep 2026 11:02:56 +0300 Subject: [PATCH 1/3] fix WriteToBigTable DirectRow --- .../beam_PostCommit_Python_Arm.json | 2 +- sdks/python/apache_beam/io/gcp/bigtableio.py | 16 +++++++++++++ .../apache_beam/io/gcp/bigtableio_test.py | 24 +++++++++++++++++++ 3 files changed, 41 insertions(+), 1 deletion(-) diff --git a/.github/trigger_files/beam_PostCommit_Python_Arm.json b/.github/trigger_files/beam_PostCommit_Python_Arm.json index 1efc8e9e4405..3f63c0c9975f 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Arm.json +++ b/.github/trigger_files/beam_PostCommit_Python_Arm.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 1 + "modification": 2 } diff --git a/sdks/python/apache_beam/io/gcp/bigtableio.py b/sdks/python/apache_beam/io/gcp/bigtableio.py index cd78deb7466c..81320a55c72c 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio.py @@ -71,6 +71,21 @@ __all__ = ['WriteToBigTable', 'ReadFromBigtable'] +def _restore_direct_row_pb_mutations(row): + # google-cloud-bigtable >= 2.44.0 stores mutations on `_mutations` and + # exposes `_get_mutation_pbs()`. Older MutationsBatcher still reads + # `_pb_mutations`. Pickling a 2.44 DirectRow onto an older worker leaves + # the instance without that attribute. + if hasattr(type(row), '_get_mutation_pbs') or hasattr(row, '_pb_mutations'): + return + mutations = getattr(row, '_mutations', None) + if mutations is None: + return + row._pb_mutations = [ + mut._to_pb() if hasattr(mut, '_to_pb') else mut for mut in mutations + ] + + class _BigTableWriteFn(beam.DoFn): """ Creates the connector can call and add_row to the batcher using each row in beam pipe line @@ -168,6 +183,7 @@ def process(self, row): # 'field1', # 'value1', # timestamp=datetime.now()) + _restore_direct_row_pb_mutations(row) self.batcher.mutate(row) def finish_bundle(self): diff --git a/sdks/python/apache_beam/io/gcp/bigtableio_test.py b/sdks/python/apache_beam/io/gcp/bigtableio_test.py index 08c33017f9c7..5026d2ef0eff 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio_test.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio_test.py @@ -377,6 +377,30 @@ def test_write_batch_error_surfaces_from_buffered_rows(self): self.assertGreater( mock_mutate.call_count, 0, 'buffered row was never flushed') + def test_write_direct_row_without_pb_mutations(self): + write_fn = bigtableio._BigTableWriteFn( + self._PROJECT_ID, + self._INSTANCE_ID, + self._TABLE_ID, + flush_count=1000, + max_row_bytes=5242880) + write_fn.table = self.table + write_fn.start_bundle() + + class Mutation: + def _to_pb(self): + return 'mutation-pb' + + class LegacyDirectRow: + def __init__(self): + self._mutations = [Mutation()] + + row = LegacyDirectRow() + with patch.object(MutationsBatcher, 'mutate') as mock_mutate: + write_fn.process(row) + mock_mutate.assert_called_once_with(row) + self.assertEqual(row._pb_mutations, ['mutation-pb']) + def test_write_close_error_is_surfaced(self): write_fn = bigtableio._BigTableWriteFn( self._PROJECT_ID, From f21cc532d35902ec5c8fb6ebf76d848d14ce4d36 Mon Sep 17 00:00:00 2001 From: aibrahiim Date: Mon, 7 Sep 2026 17:16:48 +0300 Subject: [PATCH 2/3] fix bigtable write --- .../beam_PostCommit_Python_Arm.json | 2 +- sdks/python/apache_beam/io/gcp/bigtableio.py | 20 ++++++++++++++----- .../apache_beam/io/gcp/bigtableio_test.py | 13 ++++++++++++ 3 files changed, 29 insertions(+), 6 deletions(-) diff --git a/.github/trigger_files/beam_PostCommit_Python_Arm.json b/.github/trigger_files/beam_PostCommit_Python_Arm.json index 3f63c0c9975f..bbdc3a3910ef 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Arm.json +++ b/.github/trigger_files/beam_PostCommit_Python_Arm.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 2 + "modification": 3 } diff --git a/sdks/python/apache_beam/io/gcp/bigtableio.py b/sdks/python/apache_beam/io/gcp/bigtableio.py index 81320a55c72c..f1a19e179f92 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio.py @@ -62,9 +62,11 @@ from google.cloud.bigtable import Client from google.cloud.bigtable.batcher import MutationsBatcher from google.cloud.bigtable.row import Cell + from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.row import PartialRowData except ImportError: + DirectRow = None _LOGGER.warning( 'ImportError: from google.cloud.bigtable import Client', exc_info=True) @@ -72,11 +74,10 @@ def _restore_direct_row_pb_mutations(row): - # google-cloud-bigtable >= 2.44.0 stores mutations on `_mutations` and - # exposes `_get_mutation_pbs()`. Older MutationsBatcher still reads - # `_pb_mutations`. Pickling a 2.44 DirectRow onto an older worker leaves - # the instance without that attribute. - if hasattr(type(row), '_get_mutation_pbs') or hasattr(row, '_pb_mutations'): + # google-cloud-bigtable >= 2.44.0 stores mutations on `_mutations`. + # Older MutationsBatcher reads `_pb_mutations`. Fill that attribute so + # pickle and worker batching both see protobuf mutations. + if hasattr(row, '_pb_mutations'): return mutations = getattr(row, '_mutations', None) if mutations is None: @@ -86,6 +87,15 @@ def _restore_direct_row_pb_mutations(row): ] +def _direct_row_getstate(self): + _restore_direct_row_pb_mutations(self) + return self.__dict__ + + +if DirectRow is not None: + DirectRow.__getstate__ = _direct_row_getstate + + class _BigTableWriteFn(beam.DoFn): """ Creates the connector can call and add_row to the batcher using each row in beam pipe line diff --git a/sdks/python/apache_beam/io/gcp/bigtableio_test.py b/sdks/python/apache_beam/io/gcp/bigtableio_test.py index 5026d2ef0eff..53490eff397d 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio_test.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio_test.py @@ -18,6 +18,7 @@ """Unit tests for BigTable service.""" import logging +import pickle import string import unittest import uuid @@ -401,6 +402,18 @@ def __init__(self): mock_mutate.assert_called_once_with(row) self.assertEqual(row._pb_mutations, ['mutation-pb']) + def test_direct_row_pickle_includes_pb_mutations(self): + class Mutation: + def _to_pb(self): + return 'mutation-pb' + + row = DirectRow(row_key=b'key-1') + if hasattr(row, '_pb_mutations'): + delattr(row, '_pb_mutations') + row._mutations = [Mutation()] + restored = pickle.loads(pickle.dumps(row)) + self.assertEqual(restored._pb_mutations, ['mutation-pb']) + def test_write_close_error_is_surfaced(self): write_fn = bigtableio._BigTableWriteFn( self._PROJECT_ID, From 27639601c99a3bd424fbe501946fd0b0e5d5665c Mon Sep 17 00:00:00 2001 From: aibrahiim Date: Mon, 7 Sep 2026 22:35:01 +0300 Subject: [PATCH 3/3] fix Bigtable unit test pickle error --- sdks/python/apache_beam/io/gcp/bigtableio_test.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigtableio_test.py b/sdks/python/apache_beam/io/gcp/bigtableio_test.py index 53490eff397d..305b22eedc99 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio_test.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio_test.py @@ -18,7 +18,6 @@ """Unit tests for BigTable service.""" import logging -import pickle import string import unittest import uuid @@ -402,17 +401,15 @@ def __init__(self): mock_mutate.assert_called_once_with(row) self.assertEqual(row._pb_mutations, ['mutation-pb']) - def test_direct_row_pickle_includes_pb_mutations(self): - class Mutation: - def _to_pb(self): - return 'mutation-pb' - + def test_direct_row_getstate_includes_pb_mutations(self): row = DirectRow(row_key=b'key-1') if hasattr(row, '_pb_mutations'): delattr(row, '_pb_mutations') - row._mutations = [Mutation()] - restored = pickle.loads(pickle.dumps(row)) - self.assertEqual(restored._pb_mutations, ['mutation-pb']) + mutation = MagicMock() + mutation._to_pb.return_value = 'mutation-pb' + row._mutations = [mutation] + state = bigtableio._direct_row_getstate(row) + self.assertEqual(state['_pb_mutations'], ['mutation-pb']) def test_write_close_error_is_surfaced(self): write_fn = bigtableio._BigTableWriteFn(