ST_BinEnvelope.java

package com.esri.hadoop.hive;

import java.util.EnumSet;

import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector.Category;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector.PrimitiveCategory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorUtils;

import com.esri.core.geometry.Envelope;
import com.esri.core.geometry.ogc.OGCPoint;
import com.esri.hadoop.hive.GeometryUtils.OGCType;

@Description(
		name = "ST_BinEnvelope",
		value = "_FUNC_(binsize, point) - return bin envelope for given point\n" + 
		"_FUNC_(binsize, binid) - return bin envelope for given bin ID\n")
public class ST_BinEnvelope extends GenericUDF {
	private transient boolean binSizeIsConstant;
	private transient PrimitiveObjectInspector oiBinSize;
	private transient BinUtils bins;
	
	private transient PrimitiveObjectInspector oiBinId;
	private transient HiveGeometryOIHelper binPoint;
	
	@Override
	public ObjectInspector initialize(ObjectInspector[] OIs)
			throws UDFArgumentException {
		
		if (OIs.length != 2) {
			throw new UDFArgumentException("Function takes exactly 2 arguments");
		}
		
		if (!isPrimitiveNumber(OIs[0])) {
			throw new UDFArgumentException("Argument 0 must be a number");
		}

		oiBinSize = (PrimitiveObjectInspector)OIs[0];
		
		if (isPrimitiveNumber(OIs[1])) {
			oiBinId = (PrimitiveObjectInspector)OIs[1];
		} else if (HiveGeometryOIHelper.canCreate(OIs[1])) {
			binPoint = HiveGeometryOIHelper.create(OIs, 1);
		} else {
			throw new UDFArgumentException("Argument 1 must be a number or valid geometry type");
		}
		
		return GeometryUtils.geometryTransportObjectInspector;
	}
	
	private boolean isPrimitiveNumber(ObjectInspector oi) {
		if (oi.getCategory() != Category.PRIMITIVE) {
			return false;
		}
		
		return EnumSet.of(PrimitiveCategory.DOUBLE,PrimitiveCategory.INT,PrimitiveCategory.LONG,PrimitiveCategory.SHORT, PrimitiveCategory.FLOAT)
					  .contains(((PrimitiveObjectInspector)oi).getPrimitiveCategory());
	}
	
	@Override
	public Object evaluate(DeferredObject[] args) throws HiveException {
		double binSize = PrimitiveObjectInspectorUtils.getDouble(args[0].get(), oiBinSize);
		
		
		if (!binSizeIsConstant || bins == null) {
			bins = new BinUtils(binSize);
		} 
				
		Envelope env = new Envelope();
		
		if (oiBinId != null) {
			// argument 1 is a number, attempt to get the envelope with bin ID
			if (args[1].get() == null) {
				// null bin ID argument usually means the source point was null or failed to parse
				return null; 
			}
			
			long binId = PrimitiveObjectInspectorUtils.getLong(args[1].get(), oiBinId);
			bins.queryEnvelope(binId, env);
		} else {
			// argument 1 is a geometry, attempt to get the envelope with a point
			OGCPoint point = binPoint.getPoint(args);
			
			if (point == null) {
				return null;
			}
			
			bins.queryEnvelope(point.X(), point.Y(), env);
		}

		return GeometryUtils.geometryToEsriShapeBytesWritable(env, 0, OGCType.ST_POLYGON);
	}

	@Override
	public String getDisplayString(String[] args) {
		assert(args.length == 2);
		return String.format("st_binenvelope(%s,%s)", args[0], args[1]);
	}
}