InArray.java

package com.livingsocial.hive.udf;

import com.livingsocial.hive.utils.KISSInspector;

import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.exec.UDFArgumentLengthException;
import org.apache.hadoop.hive.ql.exec.UDFArgumentTypeException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.serde.Constants;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ListObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.io.BooleanWritable;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

@Description(name = "in_array", value = "_FUNC_(x, a) - Returns true if x is in the array a")
public class InArray extends GenericUDF {
    private ListObjectInspector array_inspector;
    private KISSInspector element_inspector;

    static final Log LOG = LogFactory.getLog(InArray.class.getName());

    @Override
    public ObjectInspector initialize(ObjectInspector[] args) throws UDFArgumentException {
	if(args.length != 2 || !KISSInspector.isPrimitive(args[0]) || !KISSInspector.isList(args[1]))
	    throw new UDFArgumentException("in_array() takes two arguments: a primitive and an array");
	element_inspector = new KISSInspector(args[0]);
	array_inspector = (ListObjectInspector) args[1];
	return PrimitiveObjectInspectorFactory.writableBooleanObjectInspector;
    }

    @Override
    public BooleanWritable evaluate(DeferredObject[] arguments) throws HiveException {
	BooleanWritable found = new BooleanWritable(false);
	// if either argument is null, or if the list doesn't contain the same primitives
        if(arguments[0].get() == null || arguments[1].get() == null || !element_inspector.sameAsTypeIn(array_inspector))
	    return found;
	Object element = arguments[0].get();
	Object elist = arguments[1].get();
	for(int i=0; i<array_inspector.getListLength(elist); i++) {
	    LOG.warn("elem from array: " + element_inspector.get(array_inspector.getListElement(elist, i)));
	    LOG.warn("elem: " + element_inspector.get(element));
            Object listElem = array_inspector.getListElement(elist, i);
	    if(listElem != null && element_inspector.get(listElem).equals(element_inspector.get(element)))
		found.set(true);
	}
	return found;
    }

    @Override
    public String getDisplayString(String[] children) {
	assert (children.length == 2);
	return "in_array(" + children[0] + ", " + children[1] + ")";
    }
}