only unwrap single element unions when it does not match the typeInfo

Author: ruolin59

coral-schema/src/main/java/com/linkedin/coral/schema/avro/HiveSchemaWithPartnerVisitor.java

@@ -7,6 +7,7 @@
 
 import java.util.List;
 
+import org.apache.avro.Schema;
 import org.apache.hadoop.hive.serde2.typeinfo.ListTypeInfo;
 import org.apache.hadoop.hive.serde2.typeinfo.MapTypeInfo;
 import org.apache.hadoop.hive.serde2.typeinfo.PrimitiveTypeInfo;
@@ -57,6 +58,12 @@ public interface PartnerAccessor<P, FP> {
   @SuppressWarnings("MethodTypeParameterName")
   public static <P, FP, R, FR> R visit(TypeInfo typeInfo, P partner, HiveSchemaWithPartnerVisitor<P, FP, R, FR> visitor,
       PartnerAccessor<P, FP> accessor) {
+
+    boolean partnerWrapped = shouldUnwrapPartner(partner, typeInfo);
+    if (partnerWrapped) {
+      partner = (P) SchemaUtilities.extractSingleElementUnion((Schema) partner);
+    }
+    R resultSchema;
     switch (typeInfo.getCategory()) {
       case STRUCT:
         StructTypeInfo structTypeInfo = (StructTypeInfo) typeInfo;
@@ -69,25 +76,29 @@ public static <P, FP, R, FR> R visit(TypeInfo typeInfo, P partner, HiveSchemaWit
           R result = visit(fieldTypeInfo, fieldPartnerType, visitor, accessor);
           results.add(visitor.field(name, fieldTypeInfo, fieldPartner, result));
         }
-        return visitor.struct(structTypeInfo, partner, results);
+        resultSchema = visitor.struct(structTypeInfo, partner, results);
+        break;
 
       case LIST:
         ListTypeInfo listTypeInfo = (ListTypeInfo) typeInfo;
         TypeInfo elementTypeInfo = listTypeInfo.getListElementTypeInfo();
         P elementPartner = partner != null ? accessor.listElementPartner(partner) : null;
         R elementResult = visit(elementTypeInfo, elementPartner, visitor, accessor);
-        return visitor.list(listTypeInfo, partner, elementResult);
+        resultSchema = visitor.list(listTypeInfo, partner, elementResult);
+        break;
 
       case MAP:
         MapTypeInfo mapTypeInfo = (MapTypeInfo) typeInfo;
         P keyPartner = partner != null ? accessor.mapKeyPartner(partner) : null;
         R keyResult = visit(mapTypeInfo.getMapKeyTypeInfo(), keyPartner, visitor, accessor);
         P valuePartner = partner != null ? accessor.mapValuePartner(partner) : null;
         R valueResult = visit(mapTypeInfo.getMapValueTypeInfo(), valuePartner, visitor, accessor);
-        return visitor.map(mapTypeInfo, partner, keyResult, valueResult);
+        resultSchema = visitor.map(mapTypeInfo, partner, keyResult, valueResult);
+        break;
 
       case PRIMITIVE:
-        return visitor.primitive((PrimitiveTypeInfo) typeInfo, partner);
+        resultSchema = visitor.primitive((PrimitiveTypeInfo) typeInfo, partner);
+        break;
 
       case UNION:
         UnionTypeInfo unionTypeInfo = (UnionTypeInfo) typeInfo;
@@ -98,11 +109,19 @@ public static <P, FP, R, FR> R visit(TypeInfo typeInfo, P partner, HiveSchemaWit
           R result = visit(allAlternatives.get(i), unionObjectPartner, visitor, accessor);
           unionResults.add(result);
         }
-        return visitor.union(unionTypeInfo, partner, unionResults);
+        resultSchema = visitor.union(unionTypeInfo, partner, unionResults);
+        break;
 
       default:
         throw new UnsupportedOperationException(typeInfo + " not supported");
     }
+    
+    // Rewrap in single-element union if the partner was originally wrapped
+    if (partnerWrapped) {
+      resultSchema = (R) SchemaUtilities.wrapInSingleElementUnion((Schema) resultSchema);
+    }
+    
+    return resultSchema;
   }
 
   public R struct(StructTypeInfo struct, P partner, List<FR> fieldResults) {
@@ -128,4 +147,19 @@ public R primitive(PrimitiveTypeInfo primitive, P partner) {
   public R union(UnionTypeInfo union, P partner, List<R> results) {
     return null;
   }
+
+  /**
+   * Checks if a partner schema is a single-element union that needs unwrapping for matching.
+   * Returns true when the Avro partner has a single-element union but the Hive type does not declare a union.
+   *
+   * @param partner The partner schema (may be a single-element union)
+   * @param hiveTypeInfo The Hive type info for comparison
+   * @param <P> The partner schema type
+   * @return true if unwrapping (and later rewrapping) is needed, false otherwise
+   */
+  private static <P> boolean shouldUnwrapPartner(P partner, TypeInfo hiveTypeInfo) {
+    return partner != null && partner instanceof org.apache.avro.Schema
+        && SchemaUtilities.isSingleElementUnion((org.apache.avro.Schema) partner)
+        && !(hiveTypeInfo instanceof UnionTypeInfo);
+  }
 }

coral-schema/src/main/java/com/linkedin/coral/schema/avro/MergeHiveSchemaWithAvro.java

@@ -157,27 +157,6 @@ private static class AvroPartnerAccessor implements PartnerAccessor<Schema, Sche
 
     private static final Schema MAP_KEY = Schema.create(Schema.Type.STRING);
 
-    /**
-     * Extracts the actual type from a schema that may be a union.
-     * Handles two cases:
-     * 1. Single-element unions like [{"type":"record",...}] - extracts the single type
-     * 2. Nullable unions like ["null", "string"] - extracts the non-null type
-     *
-     * @param schema The schema to extract from
-     * @return The extracted schema
-     */
-    private static Schema extractFromUnion(Schema schema) {
-      if (schema.getType() == Schema.Type.UNION) {
-        List<Schema> types = schema.getTypes();
-        if (types.size() == 1) {
-          // Single-element union: just extract that single type
-          return types.get(0);
-        }
-      }
-      // Extract the non-null type from nullable unions like ["null", "string"]
-      return SchemaUtilities.extractIfOption(schema);
-    }
-
     @Override
     public Schema.Field fieldPartner(Schema partner, String fieldName) {
       Schema schema = SchemaUtilities.extractIfOption(partner);
@@ -201,7 +180,7 @@ public Schema mapValuePartner(Schema partner) {
       if (schema.getType() != Schema.Type.MAP) {
         return null;
       }
-      return extractFromUnion(schema.getValueType());
+      return schema.getValueType();
     }
 
     @Override
@@ -210,7 +189,7 @@ public Schema listElementPartner(Schema partner) {
       if (schema.getType() != Schema.Type.ARRAY) {
         return null;
       }
-      return extractFromUnion(schema.getElementType());
+      return schema.getElementType();
     }
 
     @Override

coral-schema/src/main/java/com/linkedin/coral/schema/avro/SchemaUtilities.java

@@ -560,6 +560,32 @@ static Schema extractIfOption(Schema schema) {
     }
   }
 
+  /**
+   * Checks if a schema is a single-element union.
+   * Single-element unions are unions with exactly one type, like [{"type":"record",...}]
+   *
+   * @param schema The schema to check
+   * @return true if the schema is a single-element union, false otherwise
+   */
+  static boolean isSingleElementUnion(Schema schema) {
+    return schema != null && schema.getType() == Schema.Type.UNION && schema.getTypes().size() == 1;
+  }
+
+  /**
+   * Extracts the actual type from a single-element union.
+   * Assumes the schema is a single-element union (caller should check with isSingleElementUnion first).
+   *
+   * @param schema The single-element union schema
+   * @return The extracted schema (the single type from the union)
+   */
+  static Schema extractSingleElementUnion(Schema schema) {
+    return schema.getTypes().get(0);
+  }
+
+  static Schema wrapInSingleElementUnion(Schema schema) {
+    return Schema.createUnion(Arrays.asList(schema));
+  }
+
   private static Schema getUnionFieldSchema(@Nonnull Schema leftSchema, @Nonnull Schema rightSchema,
       boolean strictMode) {
     Preconditions.checkNotNull(leftSchema);

coral-schema/src/test/java/com/linkedin/coral/schema/avro/MergeHiveSchemaWithAvroTests.java

@@ -241,84 +241,106 @@ public void shouldHandleUnions() {
   }
 
   @Test
-  public void shouldHandleSingleElementUnionsInArraysAndMaps() {
-    // This test verifies that single-element unions in array items and map values are properly unwrapped
+  public void shouldHandleSingleElementUnionsInAllTypes() {
+    // This test verifies that single-element unions in primitives, array items, and map values are properly unwrapped
     // and the nested field nullability is preserved during schema merging.
     // This reproduces the fix for handling avro.schema.literal with single-element unions like:
+    // - Primitives: "type": ["string"]
     // - Array items: "items": [{"type":"record",...}]
     // - Map values: "values": [{"type":"record",...}]
     // These single-element unions appear in real-world Avro schemas stored as avro.schema.literal
 
-    String hive = "struct<id:bigint,items:array<struct<fooconfiguration:struct<name:string,urlvalue:string,source:string>,"
+    // seu = single-element union
+    String hive = "struct<id:bigint,name:string,active:boolean,"
+        + "items:array<struct<fooconfiguration:struct<name:string,urlvalue:string,source:string>,"
         + "barconfiguration:struct<name:string,domain:string>>>,"
-        + "metadata:map<string,struct<category:string,priority:int>>>";
-
-    // Define an Avro schema literal similar to what would be stored in avro.schema.literal table property
-    // Note the single-element unions in array items and map values: [{"type":"record",...}]
-    String avroSchemaLiteral =
-        "{\"type\":\"record\",\"name\":\"test_complex_array_table\",\"namespace\":\"com.example.test\",\"fields\":["
-            + "{\"name\":\"id\",\"type\":[\"null\",\"long\"],\"default\":null},"
-            + "{\"name\":\"items\",\"type\":[\"null\",{\"type\":\"array\",\"items\":[{\"type\":\"record\",\"name\":\"ItemConfig\",\"namespace\":\"com.example.data\",\"fields\":["
-            + "{\"name\":\"fooConfiguration\",\"type\":[\"null\",{\"type\":\"record\",\"name\":\"FooConfiguration\",\"fields\":["
-            + "{\"name\":\"name\",\"type\":\"string\"},"
-            + "{\"name\":\"urlValue\",\"type\":\"string\"},"
-            + "{\"name\":\"source\",\"type\":\"string\"}"
-            + "]}],\"default\":null},"
-            + "{\"name\":\"barConfiguration\",\"type\":[\"null\",{\"type\":\"record\",\"name\":\"BarConfiguration\",\"fields\":["
-            + "{\"name\":\"name\",\"type\":\"string\"},"
-            + "{\"name\":\"domain\",\"type\":\"string\"}"
-            + "]}],\"default\":null}"
-            + "]}]}],\"default\":null},"
-            + "{\"name\":\"metadata\",\"type\":[\"null\",{\"type\":\"map\",\"values\":[{\"type\":\"record\",\"name\":\"MetadataValue\",\"namespace\":\"com.example.data\",\"fields\":["
-            + "{\"name\":\"category\",\"type\":\"string\"},"
-            + "{\"name\":\"priority\",\"type\":\"int\"}"
-            + "]}]}],\"default\":null}"
-            + "]}";
-
-    Schema avro = new Schema.Parser().parse(avroSchemaLiteral);
-    Schema merged = merge(hive, avro);
+        + "metadata:map<string,struct<category:string,priority:int>>,"
+        + "tags:array<string>>";
+
+    // Nested record schemas
+    Schema fooConfigSchema = struct("FooConfiguration", "doc-foo", "com.example.data",
+        required("name", Schema.Type.STRING), required("urlValue", Schema.Type.STRING),
+        required("source", Schema.Type.STRING));
+
+    Schema barConfigSchema = struct("BarConfiguration", "doc-bar", "com.example.data",
+        required("name", Schema.Type.STRING), required("domain", Schema.Type.STRING));
+
+    Schema seu_arrayItemConfigSchema = struct("ItemConfig", "doc-item", "com.example.data",
+        optional("fooConfiguration", fooConfigSchema), optional("barConfiguration", barConfigSchema));
+
+    Schema seu_mapValueMetadataSchema = struct("MetadataValue", "doc-metadata", "com.example.data",
+        required("category", Schema.Type.STRING), required("priority", Schema.Type.INT));
+
+    // Construct Avro schema with single-element unions for primitives, array items, and map values
+    Schema avro = struct("test_complex_array_table", "doc-test", "com.example.test", optional("id", Schema.Type.LONG),
+        field("name", singleElementUnion(Schema.create(Schema.Type.STRING)), null, "unknown", null),
+        field("active", singleElementUnion(Schema.create(Schema.Type.BOOLEAN)), null, false, null),
+        optional("items", array(singleElementUnion(seu_arrayItemConfigSchema))),
+        optional("metadata", map(singleElementUnion(seu_mapValueMetadataSchema))),
+        optional("tags", array(singleElementUnion(Schema.create(Schema.Type.STRING)))));
+
+    // Expected schema after merge: single-element unions should be preserved
+    // The structure of the Avro schema is maintained, including single-element unions
+    // Expected output should match the input avro schema
+    Schema expected = avro;
 
-    // Verify that single-element unions were properly handled
-    // Extract items array
-    Schema mergedItemsArray = SchemaUtilities.extractIfOption(merged.getField("items").schema());
-    Schema mergedItemConfig = mergedItemsArray.getElementType();
-
-    // The fix ensures that single-element union [ItemConfig] is unwrapped to ItemConfig
-    // Without the fix, this would fail because the union wouldn't be unwrapped
-    assertEquals(mergedItemConfig.getType(), Schema.Type.RECORD, "Array element should be a record, not a union");
-
-    // Extract fooConfiguration and verify nested field nullability is preserved
-    Schema mergedFooConfig =
-        SchemaUtilities.extractIfOption(mergedItemConfig.getField("fooConfiguration").schema());
-
-    // Nested fields should be non-nullable (required) as defined in the avro.schema.literal
-    assertEquals(mergedFooConfig.getField("name").schema().getType(), Schema.Type.STRING,
-        "name field should be non-nullable string");
-    assertEquals(mergedFooConfig.getField("urlValue").schema().getType(), Schema.Type.STRING,
-        "urlValue field should be non-nullable string");
-    assertEquals(mergedFooConfig.getField("source").schema().getType(), Schema.Type.STRING,
-        "source field should be non-nullable string");
-
-    // Verify barConfiguration nested fields
-    Schema mergedBarConfig =
-        SchemaUtilities.extractIfOption(mergedItemConfig.getField("barConfiguration").schema());
-    assertEquals(mergedBarConfig.getField("name").schema().getType(), Schema.Type.STRING,
-        "bar name field should be non-nullable string");
-    assertEquals(mergedBarConfig.getField("domain").schema().getType(), Schema.Type.STRING,
-        "domain field should be non-nullable string");
-
-    // Extract metadata map value and verify
-    // Ensures that single-element union [MetadataValue] is unwrapped to MetadataValue
-    Schema mergedMetadataMap = SchemaUtilities.extractIfOption(merged.getField("metadata").schema());
-    Schema mergedMetadataValue = mergedMetadataMap.getValueType();
-
-    assertEquals(mergedMetadataValue.getType(), Schema.Type.RECORD, "Map value should be a record, not a union");
-
-    // Fields in MetadataValue should be non-nullable as defined in avro.schema.literal
-    assertEquals(mergedMetadataValue.getField("category").schema().getType(), Schema.Type.STRING,
-        "category field should be non-nullable string");
-    assertEquals(mergedMetadataValue.getField("priority").schema().getType(), Schema.Type.INT,
-        "priority field should be non-nullable int");
+    Schema actual = merge(hive, avro);
+    
+    System.out.println("\n=== INPUT AVRO SCHEMA ===");
+    System.out.println(avro.toString(true));
+    System.out.println("\n=== EXPECTED OUTPUT SCHEMA ===");
+    System.out.println(expected.toString(true));
+    System.out.println("\n=== ACTUAL OUTPUT SCHEMA ===");
+    System.out.println(actual.toString(true));
+    System.out.println("\n=== END ===\n");
+
+    assertSchema(expected, actual);
+  }
+
+  @Test
+  public void shouldHandleSingleElementUnionsWithHiveUnionType() {
+    // This test ensures backward compatibility with Hive union-encoded-as-struct format
+    // when single-element unions are present in array items and map values.
+    // In Hive, unions are represented as `uniontype<type1,type2>` and encoded as
+    // struct<tag:int,field0:type1,field1:type2> in the schema.
+    // This test verifies that extractIfOption correctly handles single-element unions
+    // in nested structures even when the Hive schema uses uniontype format.
+
+    // seu = single-element union
+    String hive = "struct<id:bigint,"
+        + "status:uniontype<string,int>,"
+        + "items:array<uniontype<struct<value:string>>>,"
+        + "metadata:map<string,uniontype<struct<priority:int>>>>";
+
+    Schema seu_arrayItemSchema =
+        struct("Item", "doc-item", "com.example.data", required("value", Schema.Type.STRING));
+
+    Schema seu_mapValueMetadataSchema =
+        struct("MetadataInfo", "doc-metadata", "com.example.data", required("priority", Schema.Type.INT));
+
+    // Construct Avro schema with single-element unions in array items and map values
+    // Also includes a regular union for the status field (to test Hive uniontype compatibility)
+    Schema avro = struct("test_union_compat", "doc-test", "com.example.test", optional("id", Schema.Type.LONG),
+        required("status", union(Schema.Type.NULL, Schema.Type.STRING, Schema.Type.INT)),
+        optional("items", array(singleElementUnion(seu_arrayItemSchema))),
+        optional("metadata", map(singleElementUnion(seu_mapValueMetadataSchema))));
+
+    // Expected schema: single-element unions should be preserved, regular union preserved
+    // The Avro schema structure is maintained
+    // Expected output should match the input avro schema
+    Schema expected = avro;
+
+    Schema actual = merge(hive, avro);
+    
+    System.out.println("\n=== INPUT AVRO SCHEMA (HiveUnionType test) ===");
+    System.out.println(avro.toString(true));
+    System.out.println("\n=== EXPECTED OUTPUT SCHEMA (HiveUnionType test) ===");
+    System.out.println(expected.toString(true));
+    System.out.println("\n=== ACTUAL OUTPUT SCHEMA (HiveUnionType test) ===");
+    System.out.println(actual.toString(true));
+    System.out.println("\n=== END ===\n");
+
+    assertSchema(expected, actual);
   }
 
   // TODO: tests to retain schema props
@@ -331,6 +353,10 @@ private Schema union(Schema.Type... types) {
     return Schema.createUnion(Arrays.stream(types).map(Schema::create).collect(Collectors.toList()));
   }
 
+  private Schema singleElementUnion(Schema schema) {
+    return Schema.createUnion(Arrays.asList(schema));
+  }
+
   private void assertSchema(Schema expected, Schema actual) {
     assertEquals(actual.toString(true), expected.toString(true));
   }