Let coral hive2rel treat a union typeinfo as a struct typeinfo, inste… (

#86)

* Let coral hive2rel treat a union typeinfo as a struct typeinfo, instead of throwing exception

* Fix style issues

* Added unit test to test the extract_union UDF

coral-hive/src/main/java/com/linkedin/coral/hive/hive2rel/TypeConverter.java

@@ -7,6 +7,8 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 import com.google.common.base.Function;
 import com.google.common.collect.Lists;
@@ -145,9 +147,15 @@ public static RelDataType convert(StructTypeInfo structType, final RelDataTypeFa
     return dtFactory.createTypeWithNullability(rowType, true);
   }
 
+  // Mimic the StructTypeInfo conversion to convert a UnionTypeInfo to the corresponding RelDataType
   public static RelDataType convert(UnionTypeInfo unionType, RelDataTypeFactory dtFactory) {
-    // Union type is not supported in Calcite.
-    throw new RuntimeException("Union type is not supported");
+    List<RelDataType> fTypes = unionType.getAllUnionObjectTypeInfos().stream()
+        .map(typeInfo -> convert(typeInfo, dtFactory)).collect(Collectors.toList());
+    List<String> fNames = IntStream.range(0, unionType.getAllUnionObjectTypeInfos().size()).mapToObj(i -> "tag_" + i)
+        .collect(Collectors.toList());
+
+    RelDataType rowType = dtFactory.createStructType(fTypes, fNames);
+    return dtFactory.createTypeWithNullability(rowType, true);
   }
 
   public static TypeInfo convert(RelDataType rType) {

coral-hive/src/main/java/com/linkedin/coral/hive/hive2rel/functions/StaticHiveFunctionRegistry.java

@@ -315,6 +315,7 @@ public RelDataType inferReturnType(SqlOperatorBinding opBinding) {
 
     createAddUserDefinedFunction("array_contains", ReturnTypes.BOOLEAN, family(SqlTypeFamily.ARRAY, SqlTypeFamily.ANY));
     createAddUserDefinedFunction("sort_array", ARG0, ARRAY);
+    createAddUserDefinedFunction("extract_union", ARG0, or(ANY, family(SqlTypeFamily.ANY, SqlTypeFamily.INTEGER)));
 
     // LinkedIn UDFs: Dali stores mapping from UDF name to the implementing Java class as table properties
     // in the HCatalog. So, an UDF implementation may be referred by different names by different views.

coral-hive/src/test/java/com/linkedin/coral/hive/hive2rel/HiveToRelConverterTest.java

@@ -31,6 +31,7 @@
 import static com.linkedin.coral.hive.hive2rel.ToRelConverter.*;
 import static org.apache.calcite.sql.type.OperandTypes.*;
 import static org.testng.Assert.*;
+import static org.testng.Assert.assertEquals;
 
 
 public class HiveToRelConverterTest {
@@ -332,6 +333,21 @@ public void testCurrentUser() {
     assertEquals(generated, expected);
   }
 
+  @Test
+  public void testUnionExtractUDF() {
+    final String sql1 = "SELECT extract_union(foo) from union_table";
+    String generated1 = relToString(sql1);
+    final String expected1 =
+        "LogicalProject(EXPR$0=[extract_union($0)])\n" + "  LogicalTableScan(table=[[hive, default, union_table]])\n";
+    assertEquals(generated1, expected1);
+
+    final String sql2 = "SELECT extract_union(foo, 1) from union_table";
+    String generated2 = relToString(sql2);
+    final String expected2 = "LogicalProject(EXPR$0=[extract_union($0, 1)])\n"
+        + "  LogicalTableScan(table=[[hive, default, union_table]])\n";
+    assertEquals(generated2, expected2);
+  }
+
   private String relToString(String sql) {
     return RelOptUtil.toString(converter.convertSql(sql));
   }

coral-hive/src/test/java/com/linkedin/coral/hive/hive2rel/TestUtils.java

@@ -159,11 +159,15 @@ public static TestHive setupDefaultHive() throws IOException {
       if (response.getResponseCode() != 0) {
         throw new RuntimeException("Failed to setup view");
       }
+
+      driver.run(
+          "CREATE TABLE IF NOT EXISTS union_table(foo uniontype<int, double, array<string>, struct<a:int,b:string>>)");
+
       testHive.databases =
           ImmutableList.of(new TestHive.DB("test", ImmutableList.of("tableOne", "tableTwo", "tableOneView")),
               new TestHive.DB("default",
                   ImmutableList.of("bar", "complex", "foo", "foo_view", "null_check_view", "null_check_wrapper",
-                      "schema_evolve", "view_schema_evolve", "view_schema_evolve_wrapper")),
+                      "schema_evolve", "view_schema_evolve", "view_schema_evolve_wrapper", "union_table")),
               new TestHive.DB("fuzzy_union",
                   ImmutableList.of("tableA", "tableB", "tableC", "union_view", "union_view_with_more_than_two_tables",
                       "union_view_with_alias", "union_view_single_branch_evolved",

coral-spark/src/test/java/com/linkedin/coral/spark/CoralSparkTest.java

@@ -328,6 +328,17 @@ public void testSchemaPromotionView() {
     assertEquals(CoralSpark.create(relNode).getSparkSql(), targetSql);
   }
 
+  @Test
+  public void testUnionExtractUDF() {
+    RelNode relNode = TestUtils.toRelNode("SELECT extract_union(foo) from union_table");
+    String targetSql = String.join("\n", "SELECT extract_union(foo)", "FROM default.union_table");
+    assertEquals(CoralSpark.create(relNode).getSparkSql(), targetSql);
+
+    RelNode relNode2 = TestUtils.toRelNode("SELECT extract_union(foo, 2) from union_table");
+    String targetSql2 = String.join("\n", "SELECT extract_union(foo, 2)", "FROM default.union_table");
+    assertEquals(CoralSpark.create(relNode2).getSparkSql(), targetSql2);
+  }
+
   private List<String> convertToListOfUriStrings(List<URI> listOfUris) {
     List<String> listOfUriStrings = new LinkedList<>();
     for (URI uri : listOfUris) {

coral-spark/src/test/java/com/linkedin/coral/spark/TestUtils.java

@@ -181,6 +181,9 @@ public static void initializeViews() throws HiveException, MetaException {
     run(driver, String.join("\n", "",
         "CREATE VIEW IF NOT EXISTS view_schema_promotion_wrapper AS SELECT * from view_schema_promotion"));
     run(driver, String.join("\n", "", "ALTER TABLE schema_promotion CHANGE COLUMN b b array<double>"));
+
+    run(driver,
+        "CREATE TABLE IF NOT EXISTS union_table(foo uniontype<int, double, array<string>, struct<a:int,b:string>>)");
   }
 
   public static RelNode toRelNode(String db, String view) {