Coral-schema: Add primitive type promotion (#312)

linkedin · Oct 3, 2022 · d0a4708 · d0a4708
1 parent f49b008
commit d0a4708
Show file tree

Hide file tree

Showing 7 changed files with 109 additions and 0 deletions.
diff --git a/coral-schema/src/main/java/com/linkedin/coral/schema/avro/SchemaUtilities.java b/coral-schema/src/main/java/com/linkedin/coral/schema/avro/SchemaUtilities.java
@@ -585,6 +585,24 @@ private static Schema getUnionFieldSchema(@Nonnull Schema leftSchema, @Nonnull S
       if (ImmutableSet.of(FIXED, BYTES).equals(types)) {
         return Schema.create(BYTES);
       }
+      if (ImmutableSet.of(INT, LONG).equals(types)) {
+        return Schema.create(LONG);
+      }
+      if (ImmutableSet.of(INT, FLOAT).equals(types)) {
+        return Schema.create(FLOAT);
+      }
+      if (ImmutableSet.of(INT, DOUBLE).equals(types)) {
+        return Schema.create(DOUBLE);
+      }
+      if (ImmutableSet.of(LONG, FLOAT).equals(types)) {
+        return Schema.create(FLOAT);
+      }
+      if (ImmutableSet.of(LONG, DOUBLE).equals(types)) {
+        return Schema.create(DOUBLE);
+      }
+      if (ImmutableSet.of(FLOAT, DOUBLE).equals(types)) {
+        return Schema.create(DOUBLE);
+      }
     }
 
     throw new RuntimeException("Found two incompatible schemas for LogicalUnion operator. Left schema is: "

diff --git a/coral-schema/src/test/java/com/linkedin/coral/schema/avro/TestUtils.java b/coral-schema/src/test/java/com/linkedin/coral/schema/avro/TestUtils.java
@@ -100,6 +100,7 @@ private static void initializeTables() {
     String baseNestedUnionSchema = loadSchema("base-nested-union.avsc");
     String baseComplexLowercase = loadSchema("base-complex-lowercase.avsc");
     String baseComplexNullableWithDefaults = loadSchema("base-complex-nullable-with-defaults.avsc");
+    String basePrimitive = loadSchema("base-primitive.avsc");
 
     executeCreateTableQuery("default", "basecomplex", baseComplexSchema);
     executeCreateTableQuery("default", "basecomplexunioncompatible", baseComplexUnionCompatible);
@@ -113,6 +114,7 @@ private static void initializeTables() {
     executeCreateTableQuery("default", "basecomplexuniontype", baseComplexUnionTypeSchema);
     executeCreateTableQuery("default", "basenestedunion", baseNestedUnionSchema);
     executeCreateTableQuery("default", "basecomplexlowercase", baseComplexLowercase);
+    executeCreateTableQuery("default", "baseprimitive", basePrimitive);
     executeCreateTableWithPartitionQuery("default", "basecasepreservation", baseCasePreservation);
     executeCreateTableWithPartitionFieldSchemaQuery("default", "basecomplexfieldschema", baseComplexFieldSchema);
     executeCreateTableWithPartitionQuery("default", "basenestedcomplex", baseNestedComplexSchema);

diff --git a/...l-schema/src/test/java/com/linkedin/coral/schema/avro/ViewToAvroSchemaConverterTests.java b/...l-schema/src/test/java/com/linkedin/coral/schema/avro/ViewToAvroSchemaConverterTests.java
@@ -1037,5 +1037,44 @@ public void testCaseCallWithNullBranchAndComplexDataTypeBranch() {
         TestUtils.loadSchema("testCaseCallWithNullBranchAndComplexDataTypeBranch-expected.avsc"));
   }
 
+  @Test
+  public void testUnionIntAndLongPromoteToLong() {
+    String viewSql =
+        "CREATE VIEW v AS SELECT t1.int_col AS f1 FROM baseprimitive t1 UNION ALL SELECT t2.long_col AS f1 FROM baseprimitive t2";
+    TestUtils.executeCreateViewQuery("default", "v", viewSql);
+
+    ViewToAvroSchemaConverter viewToAvroSchemaConverter = ViewToAvroSchemaConverter.create(hiveMetastoreClient);
+    Schema actualSchema = viewToAvroSchemaConverter.toAvroSchema("default", "v");
+
+    Assert.assertEquals(actualSchema.toString(true),
+        TestUtils.loadSchema("testUnionIntAndLongPromoteToLong-expected.avsc"));
+  }
+
+  @Test
+  public void testUnionIntAndDoublePromoteToDouble() {
+    String viewSql =
+        "CREATE VIEW v AS SELECT t1.int_col AS f1 FROM baseprimitive t1 UNION ALL SELECT t2.double_col AS f1 FROM baseprimitive t2";
+    TestUtils.executeCreateViewQuery("default", "v", viewSql);
+
+    ViewToAvroSchemaConverter viewToAvroSchemaConverter = ViewToAvroSchemaConverter.create(hiveMetastoreClient);
+    Schema actualSchema = viewToAvroSchemaConverter.toAvroSchema("default", "v");
+
+    Assert.assertEquals(actualSchema.toString(true),
+        TestUtils.loadSchema("testUnionIntAndDoublePromoteToDouble-expected.avsc"));
+  }
+
+  @Test
+  public void testUnionFloatAndDoublePromoteToDouble() {
+    String viewSql =
+        "CREATE VIEW v AS SELECT t1.double_col AS f1 FROM baseprimitive t1 UNION ALL SELECT t2.float_col AS f1 FROM baseprimitive t2";
+    TestUtils.executeCreateViewQuery("default", "v", viewSql);
+
+    ViewToAvroSchemaConverter viewToAvroSchemaConverter = ViewToAvroSchemaConverter.create(hiveMetastoreClient);
+    Schema actualSchema = viewToAvroSchemaConverter.toAvroSchema("default", "v");
+
+    Assert.assertEquals(actualSchema.toString(true),
+        TestUtils.loadSchema("testUnionFloatAndDoublePromoteToDouble-expected.avsc"));
+  }
+
   // TODO: add more unit tests
 }
diff --git a/coral-schema/src/test/resources/base-primitive.avsc b/coral-schema/src/test/resources/base-primitive.avsc
@@ -0,0 +1,23 @@
+{
+  "type" : "record",
+  "name" : "BasePrimitive",
+  "namespace" : "coral.schema.avro.base.complex",
+  "fields" : [
+    {
+    "name" : "int_col",
+    "type" : "int"
+    },
+    {
+      "name": "long_col",
+      "type": "long"
+    },
+    {
+      "name": "float_col",
+      "type": "float"
+    },
+    {
+      "name": "double_col",
+      "type": "double"
+    }
+  ]
+}
diff --git a/coral-schema/src/test/resources/testUnionFloatAndDoublePromoteToDouble-expected.avsc b/coral-schema/src/test/resources/testUnionFloatAndDoublePromoteToDouble-expected.avsc
@@ -0,0 +1,9 @@
+{
+  "type" : "record",
+  "name" : "v",
+  "namespace" : "default.v",
+  "fields" : [ {
+    "name" : "f1",
+    "type" : "double"
+  } ]
+}
diff --git a/coral-schema/src/test/resources/testUnionIntAndDoublePromoteToDouble-expected.avsc b/coral-schema/src/test/resources/testUnionIntAndDoublePromoteToDouble-expected.avsc
@@ -0,0 +1,9 @@
+{
+  "type" : "record",
+  "name" : "v",
+  "namespace" : "default.v",
+  "fields" : [ {
+    "name" : "f1",
+    "type" : "double"
+  } ]
+}
diff --git a/coral-schema/src/test/resources/testUnionIntAndLongPromoteToLong-expected.avsc b/coral-schema/src/test/resources/testUnionIntAndLongPromoteToLong-expected.avsc
@@ -0,0 +1,9 @@
+{
+  "type" : "record",
+  "name" : "v",
+  "namespace" : "default.v",
+  "fields" : [ {
+    "name" : "f1",
+    "type" : "long"
+  } ]
+}