Added the third argument to ndcg UDAF

myui · Apr 6, 2016 · 1613902 · 1613902
1 parent f9561d6
commit 1613902
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 4 deletions.
diff --git a/core/src/main/java/hivemall/evaluation/NDCGUDAF.java b/core/src/main/java/hivemall/evaluation/NDCGUDAF.java
@@ -26,6 +26,7 @@
 import javax.annotation.Nonnull;
 
 import org.apache.hadoop.hive.ql.exec.Description;
+import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
 import org.apache.hadoop.hive.ql.exec.UDFArgumentTypeException;
 import org.apache.hadoop.hive.ql.metadata.HiveException;
 import org.apache.hadoop.hive.ql.parse.SemanticException;
@@ -43,18 +44,30 @@
 import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;
 import org.apache.hadoop.io.LongWritable;
 
-@Description(name = "ndcg", value = "_FUNC_(array rankItems, array correctItems)"
-        + " - Returns nDCG")
+@Description(
+        name = "ndcg",
+        value = "_FUNC_(array rankItems, array correctItems [, const boolean binaryResponses = true])"
+                + " - Returns nDCG")
 public final class NDCGUDAF extends AbstractGenericUDAFResolver {
 
     // prevent instantiation
     private NDCGUDAF() {}
 
     @Override
     public GenericUDAFEvaluator getEvaluator(@Nonnull TypeInfo[] typeInfo) throws SemanticException {
-        if (typeInfo.length != 2) {
-            throw new UDFArgumentTypeException(typeInfo.length - 1, "_FUNC_ takes two arguments");
+        if (typeInfo.length != 2 && typeInfo.length != 3) {
+            throw new UDFArgumentTypeException(typeInfo.length - 1,
+                "_FUNC_ takes two or three arguments");
         }
+        boolean binaryResponses = true;
+        if (typeInfo.length == 3) {
+            binaryResponses = HiveUtils.isBooleanTypeInfo(typeInfo[2]);
+            if (binaryResponses == false) {
+                throw new UDFArgumentException(
+                    "nDCG computation for Graded Responses is not supported yet");
+            }
+        }
+
         ListTypeInfo arg1type = HiveUtils.asListTypeInfo(typeInfo[0]);
         if (!HiveUtils.isPrimitiveTypeInfo(arg1type.getListElementTypeInfo())) {
             throw new UDFArgumentTypeException(0,

diff --git a/core/src/main/java/hivemall/utils/hadoop/HiveUtils.java b/core/src/main/java/hivemall/utils/hadoop/HiveUtils.java
@@ -220,6 +220,18 @@ public static boolean isNumberTypeInfo(@Nonnull TypeInfo typeInfo) {
                 return false;
         }
     }
+
+    public static boolean isBooleanTypeInfo(@Nonnull TypeInfo typeInfo) {
+        if (typeInfo.getCategory() != ObjectInspector.Category.PRIMITIVE) {
+            return false;
+        }
+        switch (((PrimitiveTypeInfo) typeInfo).getPrimitiveCategory()) {
+            case BOOLEAN:
+                return true;
+            default:
+                return false;
+        }
+    }
 
     public static boolean isIntegerTypeInfo(@Nonnull TypeInfo typeInfo) {
         if (typeInfo.getCategory() != ObjectInspector.Category.PRIMITIVE) {