Add a SqlReturnTypeInference strategy for extract_union UDF to deal w…

…ith 1/2 arg cases (#120)
linkedin · Sep 17, 2021 · 6fd4a09 · 6fd4a09
1 parent 4cc6fc0
commit 6fd4a09
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/coral-hive/src/main/java/com/linkedin/coral/hive/hive2rel/functions/HiveReturnTypes.java b/coral-hive/src/main/java/com/linkedin/coral/hive/hive2rel/functions/HiveReturnTypes.java
@@ -53,6 +53,20 @@ public RelDataType inferReturnType(SqlOperatorBinding opBinding) {
     }
   };
 
+  public static final SqlReturnTypeInference EXTRACT_UNION_FUNCTION_RETURN_STRATEGY = opBinding -> {
+    int numArgs = opBinding.getOperandCount();
+    Preconditions.checkState(numArgs == 1 || numArgs == 2);
+    // 1-arg case
+    if (numArgs == 1) {
+      return opBinding.getOperandType(0);
+    }
+    // 2-arg case
+    else {
+      int ordinal = opBinding.getOperandLiteralValue(1, Integer.class);
+      return opBinding.getOperandType(0).getFieldList().get(ordinal).getType();
+    }
+  };
+
   public static SqlReturnTypeInference arrayOfType(final SqlTypeName typeName) {
     return new SqlReturnTypeInference() {
       @Override

diff --git a/.../src/main/java/com/linkedin/coral/hive/hive2rel/functions/StaticHiveFunctionRegistry.java b/.../src/main/java/com/linkedin/coral/hive/hive2rel/functions/StaticHiveFunctionRegistry.java
@@ -345,7 +345,8 @@ 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)));
+    createAddUserDefinedFunction("extract_union", HiveReturnTypes.EXTRACT_UNION_FUNCTION_RETURN_STRATEGY,
+        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.