fix WriteToBigTable DirectRow

.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": 3
 }

sdks/python/apache_beam/io/gcp/bigtableio.py

@@ -62,15 +62,40 @@
   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)
 
 __all__ = ['WriteToBigTable', 'ReadFromBigtable']
 
 
+def _restore_direct_row_pb_mutations(row):
+  # 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:
+    return
+  row._pb_mutations = [
+      mut._to_pb() if hasattr(mut, '_to_pb') else mut for mut in mutations
+  ]
+
+
+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
@@ -168,6 +193,7 @@ def process(self, row):
     #                     'field1',
     #                     'value1',
     #                     timestamp=datetime.now())
+    _restore_direct_row_pb_mutations(row)
     self.batcher.mutate(row)
 
   def finish_bundle(self):

sdks/python/apache_beam/io/gcp/bigtableio_test.py

@@ -377,6 +377,40 @@ 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_direct_row_getstate_includes_pb_mutations(self):
+    row = DirectRow(row_key=b'key-1')
+    if hasattr(row, '_pb_mutations'):
+      delattr(row, '_pb_mutations')
+    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(
         self._PROJECT_ID,