make getNonNullType more developer friendly

Author: the-other-tim-brown

hudi-common/src/main/java/org/apache/hudi/common/schema/HoodieSchema.java

@@ -723,14 +723,14 @@ public boolean isNullable() {
   }
 
   /**
-   * If this is a union schema, returns the non-null type.
+   * If this is a union schema, returns the non-null type. Otherwise, returns this schema.
    *
-   * @return the non-null schema from a union
-   * @throws IllegalStateException if this is not a nullable union
+   * @return the non-null schema from a union or the current schema
+   * @throws IllegalStateException if the union has more than two types
    */
   public HoodieSchema getNonNullType() {
     if (type != HoodieSchemaType.UNION) {
-      throw new IllegalStateException("Cannot get non-null type from non-union schema: " + type);
+      return this;
     }
 
     List<HoodieSchema> types = getTypes();

hudi-common/src/test/java/org/apache/hudi/common/schema/TestHoodieSchema.java

@@ -317,9 +317,7 @@ public void testUnionSchemaTypeAccess() {
       stringSchema.getTypes();
     }, "Should throw exception when getting types from non-union schema");
 
-    assertThrows(IllegalStateException.class, () -> {
-      stringSchema.getNonNullType();
-    }, "Should throw exception when getting non-null type from non-union schema");
+    assertSame(stringSchema, stringSchema.getNonNullType());
   }
 
   @Test