Fix TypeError in avro_union_type_to_beam_type with complex union types

Author: RRap0so

sdks/python/apache_beam/io/avroio.py

@@ -554,7 +554,8 @@ 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:
+      if isinstance(avro_type,
+                    str) and avro_type in AVRO_PRIMITIVES_TO_BEAM_PRIMITIVES:
         return schema_pb2.FieldType(
             atomic_type=AVRO_PRIMITIVES_TO_BEAM_PRIMITIVES[avro_type],
             nullable=True)

sdks/python/apache_beam/io/avroio_test.py

@@ -206,6 +206,34 @@ 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):
+    union_type = [
+        'null',
+        '''
+      {
+        "avro.java.string": "String",
+        "type": "string"
+      }'''
+    ]
+    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)