extend + refactor unit test

Author: KevinGe00

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

@@ -32,6 +32,7 @@
 import com.linkedin.coral.common.HiveMetastoreClient;
 import com.linkedin.coral.common.HiveMscAdapter;
 import com.linkedin.coral.common.functions.FunctionReturnTypes;
+import com.linkedin.coral.hive.hive2rel.functions.OrdinalReturnTypeInferenceV2;
 import com.linkedin.coral.hive.hive2rel.functions.StaticHiveFunctionRegistry;
 
 import static org.apache.calcite.sql.type.OperandTypes.*;
@@ -82,7 +83,7 @@ public static void registerUdfs() {
             .rowOf(ImmutableList.of("isEven", "number"), ImmutableList.of(SqlTypeName.BOOLEAN, SqlTypeName.INTEGER)),
         family(SqlTypeFamily.INTEGER));
     StaticHiveFunctionRegistry.createAddUserDefinedFunction(
-        "com.linkedin.coral.hive.hive2rel.CoralTestUDFReturnSecondArg", ReturnTypes.ARG1,
+        "com.linkedin.coral.hive.hive2rel.CoralTestUDFReturnSecondArg", new OrdinalReturnTypeInferenceV2(1),
         family(SqlTypeFamily.STRING, SqlTypeFamily.ANY));
   }
 
@@ -107,7 +108,7 @@ private static void initializeTables() {
     String baseComplexNullableWithDefaults = loadSchema("base-complex-nullable-with-defaults.avsc");
     String basePrimitive = loadSchema("base-primitive.avsc");
     String baseComplexNestedStructSameName = loadSchema("base-complex-nested-struct-same-name.avsc");
-    String testNestedStructSchema = loadSchema("testNestedStructInnerField.avsc");
+    String baseComplexMixedNullabilities = loadSchema("base-complex-mixed-nullabilities.avsc");
 
     executeCreateTableQuery("default", "basecomplex", baseComplexSchema);
     executeCreateTableQuery("default", "basecomplexunioncompatible", baseComplexUnionCompatible);
@@ -129,7 +130,7 @@ private static void initializeTables() {
     executeCreateTableWithPartitionQuery("default", "basenestedcomplex", baseNestedComplexSchema);
     executeCreateTableWithPartitionQuery("default", "basecomplexnullablewithdefaults", baseComplexNullableWithDefaults);
     executeCreateTableWithPartitionQuery("default", "basecomplexnonnullable", baseComplexNonNullable);
-    executeCreateTableWithPartitionQuery("default", "nestedStructInnerField", testNestedStructSchema);
+    executeCreateTableWithPartitionQuery("default", "basecomplexmixednullabilities", baseComplexMixedNullabilities);
 
     String baseComplexSchemaWithDoc = loadSchema("docTestResources/base-complex-with-doc.avsc");
     String baseEnumSchemaWithDoc = loadSchema("docTestResources/base-enum-with-doc.avsc");

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

@@ -234,18 +234,22 @@ public void testUdfLessThanHundred() {
   }
 
   @Test
-  public void testInnerFieldNullabilityAfterApplyingUDF() {
+  public void testPreserveNullabilitiesAfterApplyingOrdinalReturnTypeUDF() {
     String viewSql = "CREATE VIEW innerfield_with_udf "
         + "tblproperties('functions' = 'ReturnInnerStuct:com.linkedin.coral.hive.hive2rel.CoralTestUDFReturnSecondArg', "
         + "              'dependencies' = 'ivy://com.linkedin:udf:1.0') " + "AS "
         + "SELECT default_innerfield_with_udf_ReturnInnerStuct('foo', innerRecord) AS innerRecord "
-        + "FROM nestedStructInnerField";
+        + "FROM basecomplexmixednullabilities";
 
     TestUtils.executeCreateViewQuery("default", "innerfield_with_udf", viewSql);
 
     ViewToAvroSchemaConverter viewToAvroSchemaConverter = ViewToAvroSchemaConverter.create(hiveMetastoreClient);
     Schema actualSchema = viewToAvroSchemaConverter.toAvroSchema("default", "innerfield_with_udf");
-    System.out.println(actualSchema);
+
+    // Expect all fields to retain their nullability after applying the UDF, CoralTestUDFReturnSecondArg, that simply
+    // returns the second argument as is
+    Assert.assertEquals(actualSchema.toString(true),
+        TestUtils.loadSchema("testPreserveNullabilitiesAfterApplyingOrdinalReturnTypeUDF-expected.avsc"));
   }
 
   @Test

coral-schema/src/test/resources/base-complex-mixed-nullabilities.avsc

@@ -0,0 +1,45 @@
+{
+  "type": "record",
+  "name": "OuterRecord",
+  "fields": [
+    {
+      "name": "innerRecord",
+      "type": {
+        "type": "record",
+        "name": "InnerRecord",
+        "fields": [
+          {
+            "name": "String_Field_Non_Nullable",
+            "type": "string"
+          },
+          {
+            "name": "String_Field_Nullable",
+            "type": [ "string", "null" ]
+          },
+          {
+            "name" : "Int_Field_Non_Nullable",
+            "type" : "int"
+          },
+          {
+            "name" : "Int_Field_Nullable",
+            "type" : [ "int", "null" ]
+          },
+          {
+            "name" : "Array_Col_Non_Nullable",
+            "type" :  {
+              "type" : "array",
+              "items" : "string"
+            }
+          },
+          {
+            "name" : "Array_Col_Nullable",
+            "type" : [ "null", {
+              "type" : "array",
+              "items" : [ "null", "string" ]
+            } ]
+          }
+        ]
+      }
+    }
+  ]
+}

coral-schema/src/test/resources/testNestedStructInnerField.avsc

@@ -0,0 +1,38 @@
+{
+  "type" : "record",
+  "name" : "innerfield_with_udf",
+  "namespace" : "default.innerfield_with_udf",
+  "fields" : [ {
+    "name" : "innerRecord",
+    "type" : {
+      "type" : "record",
+      "name" : "InnerRecord",
+      "namespace" : "default.innerfield_with_udf.innerfield_with_udf",
+      "fields" : [ {
+        "name" : "String_Field_Non_Nullable",
+        "type" : "string"
+      }, {
+        "name" : "String_Field_Nullable",
+        "type" : [ "string", "null" ]
+      }, {
+        "name" : "Int_Field_Non_Nullable",
+        "type" : "int"
+      }, {
+        "name" : "Int_Field_Nullable",
+        "type" : [ "int", "null" ]
+      }, {
+        "name" : "Array_Col_Non_Nullable",
+        "type" : {
+          "type" : "array",
+          "items" : "string"
+        }
+      }, {
+        "name" : "Array_Col_Nullable",
+        "type" : [ "null", {
+          "type" : "array",
+          "items" : [ "null", "string" ]
+        } ]
+      } ]
+    }
+  } ]
+}