Replace getElementType with getValueType for MAP in AvroGenericRecord… (

#31653)
apache · Jun 23, 2024 · 0f2e196 · 0f2e196
1 parent 18af8c8
commit 0f2e196
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 6 deletions.
diff --git a/...src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AvroGenericRecordToStorageApiProto.java b/...src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AvroGenericRecordToStorageApiProto.java
@@ -247,17 +247,15 @@ private static TableFieldSchema fieldDescriptorFromAvroField(Schema.Field field)
         break;
       case MAP:
         Schema keyType = Schema.create(Schema.Type.STRING);
-        Schema valueType = TypeWithNullability.create(schema.getElementType()).getType();
+        Schema valueType = Schema.create(schema.getValueType().getType());
         if (valueType == null) {
           throw new RuntimeException("Unexpected null element type!");
         }
         TableFieldSchema keyFieldSchema =
-            fieldDescriptorFromAvroField(
-                new Schema.Field("key", keyType, "key of the map entry", Schema.Field.NULL_VALUE));
+            fieldDescriptorFromAvroField(new Schema.Field("key", keyType, "key of the map entry"));
         TableFieldSchema valueFieldSchema =
             fieldDescriptorFromAvroField(
-                new Schema.Field(
-                    "value", valueType, "value of the map entry", Schema.Field.NULL_VALUE));
+                new Schema.Field("value", valueType, "value of the map entry"));
         builder =
             builder
                 .setType(TableFieldSchema.Type.STRUCT)
@@ -346,7 +344,7 @@ private static Object toProtoValue(
         return toProtoValue(fieldDescriptor, type.getType(), value);
       case MAP:
         Map<CharSequence, Object> map = (Map<CharSequence, Object>) value;
-        Schema valueType = TypeWithNullability.create(avroSchema.getElementType()).getType();
+        Schema valueType = Schema.create(avroSchema.getValueType().getType());
         if (valueType == null) {
           throw new RuntimeException("Unexpected null element type!");
         }

diff --git a/...test/java/org/apache/beam/sdk/io/gcp/bigquery/AvroGenericRecordToStorageApiProtoTest.java b/...test/java/org/apache/beam/sdk/io/gcp/bigquery/AvroGenericRecordToStorageApiProtoTest.java
@@ -32,6 +32,8 @@
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.util.EnumSet;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 import java.util.stream.Collectors;
@@ -266,6 +268,25 @@ enum TestEnum {
           .noDefault()
           .endRecord();
 
+  private static final Schema SCHEMA_WITH_MAP;
+
+  static {
+    SCHEMA_WITH_MAP =
+        SchemaBuilder.record("TestMap")
+            .fields()
+            .name("nested")
+            .type()
+            .optional()
+            .type(BASE_SCHEMA)
+            .name("aMap")
+            .type()
+            .map()
+            .values()
+            .stringType()
+            .mapDefault(ImmutableMap.<String, Object>builder().put("key1", "value1").build())
+            .endRecord();
+  }
+
   private static GenericRecord baseRecord;
   private static GenericRecord logicalTypesRecord;
   private static Map<String, Object> baseProtoExpectedFields;
@@ -505,4 +526,45 @@ public void testMessageFromGenericRecordLogicalTypes() throws Exception {
     assertEquals(7, msg.getAllFields().size());
     assertBaseRecord(msg, logicalTypesProtoExpectedFields);
   }
+
+  @Test
+  public void testMessageFromGenericRecordWithMap() throws Exception {
+    // Create a GenericRecord with a map field
+    Map<String, String> mapData = new HashMap<>();
+    mapData.put("key1", "value1");
+    mapData.put("key2", "value2");
+    GenericRecord recordWithMap =
+        new GenericRecordBuilder(SCHEMA_WITH_MAP)
+            .set("nested", baseRecord)
+            .set("aMap", mapData)
+            .build();
+
+    Descriptors.Descriptor descriptor =
+        TableRowToStorageApiProto.getDescriptorFromTableSchema(
+            AvroGenericRecordToStorageApiProto.protoTableSchemaFromAvroSchema(SCHEMA_WITH_MAP),
+            true,
+            false);
+    DynamicMessage msg =
+        AvroGenericRecordToStorageApiProto.messageFromGenericRecord(
+            descriptor, recordWithMap, null, -1);
+
+    assertEquals(2, msg.getAllFields().size());
+
+    Map<String, Descriptors.FieldDescriptor> fieldDescriptors =
+        descriptor.getFields().stream()
+            .collect(Collectors.toMap(Descriptors.FieldDescriptor::getName, Functions.identity()));
+    DynamicMessage nestedMsg = (DynamicMessage) msg.getField(fieldDescriptors.get("nested"));
+    assertBaseRecord(nestedMsg, baseProtoExpectedFields);
+
+    // Assert the map field
+    List<DynamicMessage> list = (List<DynamicMessage>) msg.getField(fieldDescriptors.get("amap"));
+    // Convert the list of DynamicMessages back to a map
+    Map<String, String> actualMap = new HashMap<>();
+    for (DynamicMessage entry : list) {
+      String key = (String) entry.getField(entry.getDescriptorForType().findFieldByName("key"));
+      String value = (String) entry.getField(entry.getDescriptorForType().findFieldByName("value"));
+      actualMap.put(key, value);
+    }
+    assertEquals(mapData, actualMap);
+  }
 }