content: package com.manning.hip.ch10;
  
  import com.maxmind.geoip.LookupService;
  import org.apache.hadoop.hive.ql.exec.Description;
  import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
  import org.apache.hadoop.hive.ql.exec.UDFArgumentLengthException;
  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.ObjectInspectorConverters;
  import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
  import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
  import org.apache.hadoop.io.Text;
  
  import java.io.IOException;
  import java.net.URL;
  
  @Description(
      name = "country",
      value = "_FUNC_(ip, geolocfile) - Returns the geolocated country code " +
      " for the IP"
  )
  public class GeolocUDF extends GenericUDF {
    private LookupService geoloc;
    private ObjectInspectorConverters.Converter[] converters;
  
    @Override
    public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException {
      if (arguments.length != 2) {
        throw new UDFArgumentLengthException(
            "The function COUNTRY(ip, geolocfile) takes exactly 2 arguments.");
      }
  
      converters = new ObjectInspectorConverters.Converter[arguments.length];
      for (int i = 0; i < arguments.length; i++) {
        converters[i] = ObjectInspectorConverters.getConverter(arguments[i],
            PrimitiveObjectInspectorFactory.javaStringObjectInspector);
      }
  
      return PrimitiveObjectInspectorFactory
          .getPrimitiveJavaObjectInspector(PrimitiveObjectInspector.PrimitiveCategory.STRING);
    }
  
    @Override
    public Object evaluate(GenericUDF.DeferredObject[] arguments) throws HiveException {
      assert (arguments.length == 2);
  
      if (arguments[0].get() == null || arguments[1].get() == null) {
        return null;
      }
  
      String ip = (String) converters[0].convert(arguments[0].get());
      String filename = (String) converters[1].convert(arguments[1].get());
  
      return lookup(ip, filename);
    }
  
    protected String lookup(String ip, String filename) throws HiveException {
      try {
        if (geoloc == null) {
          URL u = getClass().getClassLoader().getResource(filename);
          if (u == null) {
            throw new HiveException("Couldn't find geolocation file '" + filename + "'");
          }
          geoloc =
              new LookupService(u.getFile(), LookupService.GEOIP_MEMORY_CACHE);
        }
  
        String countryCode = geoloc.getCountry(ip).getCode();
  
        if ("--".equals(countryCode)) {
          return null;
        }
  
        return countryCode;
      } catch (IOException e) {
        throw new HiveException("Caught IO exception", e);
      }
    }
  
    @Override
    public String getDisplayString(String[] children) {
      assert (children.length == 2);
      return "country(" + children[0] + ", " + children[1] + ")";
    }
  }
  