fix bigtable schema issue

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableReadSchemaTransformProvider.java

@@ -149,7 +149,7 @@ public abstract static class Builder {
 
       public abstract Builder setProjectId(String projectId);
 
-      public abstract Builder setFlatten(Boolean flatten);
+      public abstract Builder setFlatten(@Nullable Boolean flatten);
 
       /** Builds a {@link BigtableReadSchemaTransformConfiguration} instance. */
       public abstract BigtableReadSchemaTransformConfiguration build();

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableSimpleWriteSchemaTransformProviderIT.java

@@ -36,6 +36,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.UUID;
 import java.util.stream.Collectors;
 import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
 import org.apache.beam.sdk.io.gcp.bigtable.BigtableWriteSchemaTransformProvider.BigtableWriteSchemaTransformConfiguration;
@@ -63,7 +64,9 @@ public class BigtableSimpleWriteSchemaTransformProviderIT {
   private BigtableTableAdminClient tableAdminClient;
   private BigtableDataClient dataClient;
   private String tableId =
-      String.format("BigtableWriteIT-%tF-%<tH-%<tM-%<tS-%<tL", LocalDateTime.now(ZoneId.of("UTC")));
+      String.format(
+          "BTSimpleWriteIT-%tF-%<tH-%<tM-%<tS-%<tL-%s",
+          LocalDateTime.now(ZoneId.of("UTC")), UUID.randomUUID().toString().substring(0, 8));
   private String projectId;
   private String instanceId;
   private PTransform<PCollectionRowTuple, PCollectionRowTuple> writeTransform;

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteIT.java

@@ -39,6 +39,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
+import java.util.UUID;
 import java.util.stream.Collectors;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.PipelineResult;
@@ -79,7 +80,9 @@ public class BigtableWriteIT implements Serializable {
   private static BigtableDataClient client;
   private static BigtableTableAdminClient tableAdminClient;
   private final String tableId =
-      String.format("BigtableWriteIT-%tF-%<tH-%<tM-%<tS-%<tL", LocalDateTime.now(ZoneId.of("UTC")));
+      String.format(
+          "BigtableWriteIT-%tF-%<tH-%<tM-%<tS-%<tL-%s",
+          LocalDateTime.now(ZoneId.of("UTC")), UUID.randomUUID().toString().substring(0, 8));
 
   private String project;
 

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteSchemaTransformProviderIT.java

@@ -36,6 +36,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import java.util.UUID;
 import java.util.stream.Collectors;
 import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
 import org.apache.beam.sdk.io.gcp.bigtable.BigtableWriteSchemaTransformProvider.BigtableWriteSchemaTransformConfiguration;
@@ -64,7 +65,9 @@ public class BigtableWriteSchemaTransformProviderIT {
   private BigtableTableAdminClient tableAdminClient;
   private BigtableDataClient dataClient;
   private String tableId =
-      String.format("BigtableWriteIT-%tF-%<tH-%<tM-%<tS-%<tL", LocalDateTime.now(ZoneId.of("UTC")));
+      String.format(
+          "BTWriteSchemaIT-%tF-%<tH-%<tM-%<tS-%<tL-%s",
+          LocalDateTime.now(ZoneId.of("UTC")), UUID.randomUUID().toString().substring(0, 8));
   private String projectId;
   private String instanceId;
   private PTransform<PCollectionRowTuple, PCollectionRowTuple> writeTransform;

sdks/python/apache_beam/yaml/tests/bigtable.yaml

@@ -30,6 +30,9 @@ fixtures:
   # Tests for BigTable YAML IO
 
 pipelines:
+  # Pipeline 1: Write test data (SetCell mutations) to Bigtable, converting
+  # YAML string fields (key, column_qualifier, value) to UTF-8 bytes expected
+  # by WriteToBigTable.
   - pipeline:
       type: chain
       transforms:
@@ -83,6 +86,10 @@ pipelines:
             project: 'apache-beam-testing'
             instance: "{BT_INSTANCE}"
             table: 'test-table'
+
+  # Pipeline 2: Read from Bigtable with flatten=True (one output row per column
+  # qualifier), decode byte fields back to UTF-8 strings, and verify the
+  # flattened rows with AssertEqual.
   - pipeline:
       type: chain
       transforms:
@@ -91,6 +98,7 @@ pipelines:
             project: 'apache-beam-testing'
             instance: "{BT_INSTANCE}"
             table: 'test-table'
+            flatten: True
         - type: MapToFields
           config:
             language: python
@@ -130,6 +138,10 @@ pipelines:
                     timestamp_micros: 1000 } ] }
         - type: LogForTesting
 
+  # Pipeline 3: Read from Bigtable with flatten=False (one output row per
+  # Bigtable row key with nested column_families map), decode the key and nested
+  # cell values from bytes to UTF-8 strings, and verify the nested structure
+  # with AssertEqual (Issue #35790).
   - pipeline:
       type: chain
       transforms:
@@ -145,33 +157,41 @@ pipelines:
             fields:
               key:
                 callable: |
-                  def convert_to_bytes(row):
-                    return row.key.decode("utf-8") if "key" in row._fields else None
+                  def convert_to_string(row):
+                    k = getattr(row, 'key', None)
+                    return k.decode("utf-8") if hasattr(k, 'decode') else k
 
               column_families:
-                column_families
-#        TODO: issue  #35790, once fixed we can uncomment this assert
-#        - type: AssertEqual
-#          config:
-#            elements:
-#              - {key: 'row1',
-#                # Use explicit map syntax to match the actual output
-#                 column_families: {
-#                   cf1: {
-#                     cq1: [
-#                       { value: "value1", timestamp_micros: 5000 }
-#                     ],
-#                     cq2: [
-#                       { value: "value2", timestamp_micros: 1000 }
-#                     ]
-#                   }
-#                 }
-#              }
-        #                - {'key': 'row1',
-        #                   column_families: {cf1: {cq2:
-        #                                             [BeamSchema_3281a0ae_fe85_474b_9030_86fbed58833a(value=b'value2', timestamp_micros=1000)], 'cq1': [BeamSchema_3281a0ae_fe85_474b_9030_86fbed58833a(value=b'value1', timestamp_micros=5000)]}}}
-
-
-#        - type: LogForTesting
+                callable: |
+                  def convert_cells_to_string(row):
+                    cf = getattr(row, 'column_families', None)
+                    if not cf:
+                      return None
+                    return {
+                      fam: {
+                        col: [
+                          beam.Row(value=c.value.decode("utf-8") if hasattr(c.value, 'decode') else c.value,
+                                   timestamp_micros=c.timestamp_micros)
+                          for c in cells
+                        ]
+                        for col, cells in cols.items()
+                      }
+                      for fam, cols in cf.items()
+                    }
+        - type: AssertEqual
+          config:
+            elements:
+              - key: 'row1'
+                # Use explicit map syntax to match the actual output
+                column_families: {
+                  cf1: {
+                    cq1: [
+                      { value: "value1", timestamp_micros: 5000 }
+                    ],
+                    cq2: [
+                      { value: "value2", timestamp_micros: 1000 }
+                    ]
+                  }
+                }
 
 

sdks/python/apache_beam/yaml/yaml_provider.py

@@ -759,6 +759,23 @@ def dicts_to_rows(o):
     return o
 
 
+def to_dict(value):
+  """Recursively converts Row, NamedTuple, or Mapping objects to dicts, omitting
+  fields with None values."""
+  if value is None:
+    return None
+  if hasattr(value, '_asdict'):
+    return {k: to_dict(v) for k, v in value._asdict().items() if v is not None}
+  elif hasattr(value, 'as_dict'):
+    return {k: to_dict(v) for k, v in value.as_dict().items() if v is not None}
+  elif isinstance(value, (list, tuple)):
+    return [to_dict(v) for v in value]
+  elif isinstance(value, Mapping):
+    return {k: to_dict(v) for k, v in value.items() if v is not None}
+  else:
+    return value
+
+
 def _unify_element_with_schema(element, target_schema):
   """Convert an element to match the target schema, preserving existing
     fields only."""
@@ -828,11 +845,6 @@ def __init__(self, elements: Iterable[Any]):
       self._elements = elements
 
     def expand(self, pcoll):
-      def to_dict(row):
-        # filter None when comparing
-        temp_dict = {k: v for k, v in row._asdict().items() if v is not None}
-        return dict(temp_dict.items())
-
       return assert_that(
           pcoll | beam.Map(to_dict),
           equal_to([to_dict(e) for e in dicts_to_rows(self._elements)]))

sdks/python/apache_beam/yaml/yaml_provider_unit_test.py

@@ -377,3 +377,47 @@ def test_create_mixed_types(self):
               [('a', None), ('element', 1)],
               [('a', 2), ('element', None)],
           ]))
+
+
+class YamlProvidersAssertEqualTest(unittest.TestCase):
+  def test_assert_equal_nested_mapping(self):
+    # Issue #35790: elements with nested dictionaries / MapFields
+    with beam.Pipeline() as p:
+      input_data = [
+          beam.Row(
+              key='row1',
+              column_families={
+                  'cf1': {
+                      'cq1': [beam.Row(value='value1', timestamp_micros=5000)],
+                      'cq2': [beam.Row(value='value2', timestamp_micros=1000)]
+                  }
+              })
+      ]
+      pcoll = p | beam.Create(input_data)
+      _ = pcoll | YamlProviders.AssertEqual(
+          elements=[{
+              'key': 'row1',
+              'column_families': {
+                  'cf1': {
+                      'cq1': [{
+                          'value': 'value1', 'timestamp_micros': 5000
+                      }],
+                      'cq2': [{
+                          'value': 'value2', 'timestamp_micros': 1000
+                      }]
+                  }
+              }
+          }])
+
+  def test_assert_equal_nested_rows(self):
+    with beam.Pipeline() as p:
+      input_data = [beam.Row(key='row1', nested=beam.Row(sub=beam.Row(val=42)))]
+      pcoll = p | beam.Create(input_data)
+      _ = pcoll | YamlProviders.AssertEqual(
+          elements=[{
+              'key': 'row1', 'nested': {
+                  'sub': {
+                      'val': 42
+                  }
+              }
+          }])

sdks/python/apache_beam/yaml/yaml_testing.py

@@ -370,7 +370,7 @@ def __init__(self, elements, recording_id):
   def expand(self, pcoll):
     # Convert elements to rows outside the matcher to avoid capturing
     # any grpc channels that might be created during the conversion
-    expected_rows = yaml_provider.dicts_to_rows(self._elements)
+    expected_rows = [yaml_provider.to_dict(e) for e in self._elements]
     recording_id = self._recording_id
 
     # Create a serializable matcher function that doesn't capture
@@ -392,8 +392,7 @@ def __call__(self, actual):
             raise
 
     matcher = SerializableMatcher(expected_rows, recording_id)
-    return assert_that(
-        pcoll | beam.Map(lambda row: beam.Row(**row._asdict())), matcher)
+    return assert_that(pcoll | beam.Map(yaml_provider.to_dict), matcher)
 
 
 def create_test(
@@ -548,10 +547,7 @@ def _composite_key_to_nested(
 
 
 def _try_row_as_dict(row):
-  try:
-    return row._asdict()
-  except AttributeError:
-    return row
+  return yaml_provider.to_dict(row)
 
 
 # Linter: No need for unittest.main here.