Fix TypeError in avro_union_type_to_beam_type with complex union types

Author: RRap0so

sdks/python/apache_beam/io/avroio.py

@@ -554,11 +554,11 @@ def avro_union_type_to_beam_type(union_type: List) -> schema_pb2.FieldType:
   """
   if len(union_type) == 2 and "null" in union_type:
     for avro_type in union_type:
-      if avro_type in AVRO_PRIMITIVES_TO_BEAM_PRIMITIVES:
-        return schema_pb2.FieldType(
-            atomic_type=AVRO_PRIMITIVES_TO_BEAM_PRIMITIVES[avro_type],
-            nullable=True)
-    return schemas.typing_to_runner_api(Any)
+      if avro_type != "null":
+        beam_type = avro_type_to_beam_type(avro_type)
+        if beam_type.WhichOneof("type_info") == "atomic_type":
+          return schema_pb2.FieldType(
+              atomic_type=beam_type.atomic_type, nullable=True)
   return schemas.typing_to_runner_api(Any)
 
 

sdks/python/apache_beam/io/avroio_test.py

@@ -206,6 +206,43 @@ def test_avro_union_type_to_beam_type_with_string_long(self):
     expected_beam_type = schemas.typing_to_runner_api(Any)
     hc.assert_that(beam_type, hc.equal_to(expected_beam_type))
 
+  def test_avro_union_type_to_beam_type_with_record_and_null(self):
+    record_type = {
+        'type': 'record',
+        'name': 'TestRecord',
+        'fields': [{
+            'name': 'field1', 'type': 'string'
+        }, {
+            'name': 'field2', 'type': 'int'
+        }]
+    }
+    union_type = [record_type, 'null']
+    beam_type = avro_union_type_to_beam_type(union_type)
+    expected_beam_type = schemas.typing_to_runner_api(Any)
+    hc.assert_that(beam_type, hc.equal_to(expected_beam_type))
+
+  def test_avro_union_type_to_beam_type_with_nullable_annotated_string(self):
+    annotated_string_type = {"avro.java.string": "String", "type": "string"}
+    union_type = ['null', annotated_string_type]
+
+    beam_type = avro_union_type_to_beam_type(union_type)
+
+    expected_beam_type = schema_pb2.FieldType(
+        atomic_type=schema_pb2.STRING, nullable=True)
+    hc.assert_that(beam_type, hc.equal_to(expected_beam_type))
+
+  def test_avro_union_type_to_beam_type_with_only_null(self):
+    union_type = ['null']
+    beam_type = avro_union_type_to_beam_type(union_type)
+    expected_beam_type = schemas.typing_to_runner_api(Any)
+    hc.assert_that(beam_type, hc.equal_to(expected_beam_type))
+
+  def test_avro_union_type_to_beam_type_with_multiple_types(self):
+    union_type = ['null', 'string', 'int']
+    beam_type = avro_union_type_to_beam_type(union_type)
+    expected_beam_type = schemas.typing_to_runner_api(Any)
+    hc.assert_that(beam_type, hc.equal_to(expected_beam_type))
+
   def test_avro_schema_to_beam_and_back(self):
     avro_schema = fastavro.parse_schema(json.loads(self.SCHEMA_STRING))
     beam_schema = avro_schema_to_beam_schema(avro_schema)