From 6fd4a094c7c8e585a769d0724e75f10943770847 Mon Sep 17 00:00:00 2001
From: "Limian (Raymond) Zhang" <razhang@linkedin.com>
Date: Fri, 17 Sep 2021 07:30:36 -0700
Subject: [PATCH] Add a SqlReturnTypeInference strategy for extract_union UDF
 to deal with 1/2 arg cases (#120)

---
 .../hive/hive2rel/functions/HiveReturnTypes.java   | 14 ++++++++++++++
 .../functions/StaticHiveFunctionRegistry.java      |  3 ++-
 2 files changed, 16 insertions(+), 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
index 25d32f68d..d693caba9 100644
--- 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/coral-hive/src/main/java/com/linkedin/coral/hive/hive2rel/functions/StaticHiveFunctionRegistry.java b/coral-hive/src/main/java/com/linkedin/coral/hive/hive2rel/functions/StaticHiveFunctionRegistry.java
index 06d61bd19..caf1c08dd 100644
--- a/coral-hive/src/main/java/com/linkedin/coral/hive/hive2rel/functions/StaticHiveFunctionRegistry.java
+++ b/coral-hive/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.