[GOBBLIN-496] Support nullable union when getting field schema

Closes #2367 from jack-moseley/nullable-union
apache · May 18, 2018 · c47f43d · c47f43d
1 parent 01302a6
commit c47f43d
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/gobblin-utility/src/main/java/org/apache/gobblin/util/AvroUtils.java b/gobblin-utility/src/main/java/org/apache/gobblin/util/AvroUtils.java
@@ -57,6 +57,7 @@
 import org.apache.hadoop.fs.PathFilter;
 import org.apache.hadoop.fs.permission.FsAction;
 import org.apache.hadoop.fs.permission.FsPermission;
+import org.apache.hadoop.hive.serde2.avro.AvroSerdeUtils;
 import org.codehaus.jackson.JsonNode;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -126,6 +127,9 @@ private static Optional<Schema> getFieldSchemaHelper(Schema schema, List<String>
     }
     switch (schema.getType()) {
       case UNION:
+        if (AvroSerdeUtils.isNullableType(schema)) {
+          return AvroUtils.getFieldSchemaHelper(AvroSerdeUtils.getOtherTypeFromNullableType(schema), pathList, field);
+        }
         throw new AvroRuntimeException("Union of complex types cannot be handled : " + schema);
       case MAP:
         if ((field + 1) == pathList.size()) {

diff --git a/gobblin-utility/src/test/java/org/apache/gobblin/util/AvroUtilsTest.java b/gobblin-utility/src/test/java/org/apache/gobblin/util/AvroUtilsTest.java
@@ -279,4 +279,24 @@ public void testComplexTypesInUnionNotSupported()
 
     AvroUtils.getFieldSchema(record.getSchema(), TEST_LOCATION);
   }
+
+  @Test
+  public void testUnionWithNull() {
+    Schema nestedRecord = SchemaBuilder.record("nested").fields().requiredDouble("double")
+        .requiredString("string").endRecord();
+    Schema union = SchemaBuilder.unionOf().nullType().and().type(nestedRecord).endUnion();
+    Schema schema = SchemaBuilder.record("record").fields().name("union").type(union).noDefault().endRecord();
+
+    Schema doubleSchema = AvroUtils.getFieldSchema(schema, "union.double").get();
+    Assert.assertEquals(doubleSchema.getType(), Schema.Type.DOUBLE);
+
+    GenericRecord nested = new GenericData.Record(nestedRecord);
+    nested.put("double", 10);
+    nested.put("string", "testString");
+    GenericRecord record = new GenericData.Record(schema);
+    record.put("union", nested);
+
+    String stringValue = AvroUtils.getFieldValue(record, "union.string").get().toString();
+    Assert.assertEquals(stringValue, "testString");
+  }
 }