package brickhouse.udf.dcache;
/**
* Copyright 2012 Klout, Inc
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
**/
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.UDFType;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.serde2.SerDe;
import org.apache.hadoop.hive.serde2.SerDeException;
import org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe;
import org.apache.hadoop.hive.serde2.objectinspector.ConstantObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.MapObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector.Category;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorUtils;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.StructField;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.StringObjectInspector;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoFactory;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoUtils;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.Writable;
import org.apache.log4j.Level;
import org.apache.log4j.Logger;
import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
/**
* UDF to access a distributed map file
* <p/>
* Assumes the file is a tab-separated file of name-value pairs,
* which has been placed in distributed cache using the "add file" command
* <p/>
* Example
* <p/>
* INSERT OVERWRITE LOCAL DIRECTORY mymap select key,value from my_map_table;
* ADD FILE mymap;
* <p/>
* select key, val* distributed_map( key, 'mymap') from the_table;
* <p/>
* <p/>
* If one argument is passed in, it is assumed to be a filename, containing
* a map of type map<string,double>, and the entire map is returned.
* <p/>
* If two arguments are passed in, it is either filename, and a string specifying the
* type of the map ( i.e distributed_map('mymap','map<string,bigint>'); ) and returns
* the entire map, or it is the key and the filename ( ie distributed_map( key, 'mymap'),
* and only the key's value is returned.
* <p/>
* If there are three arguments passed in, it is assumed to be the key, the filename, and the
* maptype, (i.e distributed_map( key, 'mymap', 'map<string,bigint>') )
*/
@UDFType(deterministic = false)
public class DistributedMapUDF extends GenericUDF {
private static final Logger LOG = Logger.getLogger(DistributedMapUDF.class);
private static HashMap<String, HashMap<Object, Object>> localMapMap = new HashMap<String, HashMap<Object, Object>>();
private StringObjectInspector fileNameInspector;
private PrimitiveObjectInspector keyInspector;
private TypeInfo keyType;
private TypeInfo valType;
private LazySimpleSerDe serde;
private LazySimpleSerDe getLineSerde() throws SerDeException {
if (serde == null) {
Logger.getLogger(LazySimpleSerDe.class).setLevel(Level.DEBUG);
serde = new LazySimpleSerDe();
Configuration job = new Configuration();
Properties tbl = new Properties();
tbl.setProperty("columns", "key,value");
tbl.setProperty("columns.types", keyType.getTypeName() + "," + valType.getTypeName());
serde.initialize(job, tbl);
}
return serde;
}
private void addValues(HashMap<Object, Object> map, String mapFilename) throws IOException, SerDeException {
if (!mapFilename.endsWith("crc")) {
File mapFile = new File(mapFilename);
if (mapFile.isDirectory()) {
String[] subFiles = mapFile.list();
for (String subFile : subFiles) {
LOG.info("Checking recursively " + subFile);
addValues(map, mapFilename + "/" + subFile);
}
} else {
BufferedReader reader = new BufferedReader(new InputStreamReader(new FileInputStream(mapFile)));
SerDe lazy = getLineSerde();
StructObjectInspector lineInsp = (StructObjectInspector) lazy.getObjectInspector();
StructField keyRef = lineInsp.getStructFieldRef("key");
StructField valueRef = lineInsp.getStructFieldRef("value");
String line;
while ((line = reader.readLine()) != null) {
Writable lineText = new Text(line);
Object lineObj = lazy.deserialize(lineText);
List<Object> objList = lineInsp.getStructFieldsDataAsList(lineObj);
Object key = ((PrimitiveObjectInspector) keyRef.getFieldObjectInspector()).getPrimitiveJavaObject(objList.get(0));
Object val = ((PrimitiveObjectInspector) valueRef.getFieldObjectInspector()).getPrimitiveJavaObject(objList.get(1));
map.put(key, val);
}
}
} else {
LOG.info(" Ignoring CRC file " + mapFilename);
}
}
private Map<Object, Object> getLocalMap(String mapFileName) {
HashMap<Object, Object> map = localMapMap.get(mapFileName);
if (map == null) {
try {
File localDir = new File(".");
String[] files = localDir.list();
for (String file : files) {
LOG.info(" In current dir is " + file);
File checkFile = new File(file);
if (checkFile.isDirectory()) {
LOG.info(" FILE " + file + " is a directory");
}
}
map = new HashMap<Object, Object>();
addValues(map, mapFileName);
localMapMap.put(mapFileName, map);
} catch (IOException ioExc) {
ioExc.printStackTrace();
throw new RuntimeException(ioExc);
} catch (SerDeException serdeExc) {
throw new RuntimeException(serdeExc);
}
}
return map;
}
@Override
public Object evaluate(DeferredObject[] arg0) throws HiveException {
//// if keyInspector has been set
if (this.keyInspector != null) {
Object key = keyInspector.getPrimitiveJavaObject(arg0[0].get());
String mapFileName = this.fileNameInspector.getPrimitiveJavaObject(arg0[1].get());
Map<Object, Object> map = getLocalMap(mapFileName);
return map.get(key);
} else {
Object mapFNameObj;
if (arg0.length == 1) {
mapFNameObj = arg0[0].get();
} else {
mapFNameObj = arg0[1].get();
}
String mapFileName = this.fileNameInspector.getPrimitiveJavaObject(mapFNameObj);
Map<Object, Object> map = getLocalMap(mapFileName);
return map;
}
}
@Override
public String getDisplayString(String[] arg0) {
return "distributed_map()";
}
public String usage(String err) {
return " Distributed Map -- Case " + err;
}
private MapObjectInspector getMapType(String typeStr) throws UDFArgumentException, IllegalArgumentException {
try {
TypeInfo hiveType = TypeInfoUtils.getTypeInfoFromTypeString(typeStr);
if (hiveType.getCategory() != Category.MAP) {
throw new UDFArgumentException(usage("Type is not map"));
}
MapObjectInspector mapInsp = (MapObjectInspector) TypeInfoUtils.getStandardJavaObjectInspectorFromTypeInfo(hiveType);
if (mapInsp.getMapKeyObjectInspector().getCategory() != Category.PRIMITIVE) {
throw new UDFArgumentException(usage("Key is not primitive"));
}
return mapInsp;
} catch (IllegalArgumentException badTypeStr) {
throw new UDFArgumentException(usage("String is not type"));
}
}
/**
* Either one, two or three values can be passed in.
* If one argument is passed it, it is implied that the
* return value is a map<string,double>. If three arguments
* are passed in, then it is implied the arguments are the
* map key, the map filename, and the value type.
* <p/>
* If two arguments are passed in, it is implied that either)
* a map key, and a filename are being passed in,
* or a filename, and a map return type are being passed in.
*/
@Override
public ObjectInspector initialize(ObjectInspector[] arg0)
throws UDFArgumentException {
if (arg0.length == 0 || arg0.length > 3)
throw new UDFArgumentException(usage("Between 1 and 3 arguments"));
switch (arg0.length) {
case 1:
//// filename
if (!(arg0[0] instanceof ConstantObjectInspector)
|| !(arg0[0] instanceof StringObjectInspector)) {
throw new UDFArgumentException(usage(" 1 arguments is always name of directory"));
}
fileNameInspector = (StringObjectInspector) arg0[0];
keyType = TypeInfoFactory.stringTypeInfo;
valType = TypeInfoFactory.doubleTypeInfo;
return ObjectInspectorFactory.getStandardMapObjectInspector(PrimitiveObjectInspectorFactory.javaStringObjectInspector,
PrimitiveObjectInspectorFactory.javaDoubleObjectInspector);
case 2:
//// either key, filename ...
/// or filename, maptype
if (!(arg0[1] instanceof ConstantObjectInspector)
|| !(arg0[1] instanceof StringObjectInspector)
|| !(arg0[0] instanceof PrimitiveObjectInspector)) {
throw new UDFArgumentException(usage("2 arguments is eiter key and filename, or a filename and maptype"));
}
ConstantObjectInspector mapType = (ConstantObjectInspector) arg0[1];
String typeStr = mapType.getWritableConstantValue().toString();
try {
//// able to parse map type ...
MapObjectInspector mapInsp = getMapType(typeStr);
keyType = TypeInfoFactory.getPrimitiveTypeInfo(mapInsp.getMapKeyObjectInspector().getTypeName());
valType = TypeInfoFactory.getPrimitiveTypeInfo(mapInsp.getMapValueObjectInspector().getTypeName());
fileNameInspector = (StringObjectInspector) arg0[0];
return ObjectInspectorUtils.getStandardObjectInspector(mapInsp);
} catch (UDFArgumentException checkMapType) {
/// Assume that it is key, filename
this.keyInspector = (PrimitiveObjectInspector) arg0[0];
keyType = TypeInfoFactory.getPrimitiveTypeInfo(keyInspector.getTypeName());
valType = TypeInfoFactory.doubleTypeInfo;
this.fileNameInspector = (StringObjectInspector) arg0[1];
//// Default case is string, double
return PrimitiveObjectInspectorFactory.javaDoubleObjectInspector;
}
case 3:
//// key , filename , maptype
if (!(arg0[1] instanceof ConstantObjectInspector)
|| !(arg0[1] instanceof StringObjectInspector)
|| !(arg0[2] instanceof ConstantObjectInspector)
|| !(arg0[2] instanceof StringObjectInspector)
|| !(arg0[0] instanceof PrimitiveObjectInspector)) {
throw new UDFArgumentException(usage("3 arguments are key,filename and maptype"));
}
fileNameInspector = (StringObjectInspector) arg0[1];
ConstantObjectInspector mapType3 = (ConstantObjectInspector) arg0[2];
String typeStr3 = mapType3.getWritableConstantValue().toString();
MapObjectInspector mapInspect = this.getMapType(typeStr3);
keyInspector = (PrimitiveObjectInspector) arg0[0];
if (keyInspector.getPrimitiveCategory() !=
((PrimitiveObjectInspector) mapInspect.getMapKeyObjectInspector()).getPrimitiveCategory()) {
throw new UDFArgumentException(usage("Key must be primitive"));
}
keyType = TypeInfoFactory.getPrimitiveTypeInfo(keyInspector.getTypeName());
ObjectInspector valInspector = ObjectInspectorUtils.getStandardObjectInspector(mapInspect.getMapValueObjectInspector());
/// XXX Can we have non primitives for the values ????
valType = TypeInfoFactory.getPrimitiveTypeInfo(valInspector.getTypeName());
return valInspector;
}
return null;
}
}