package brickhouse.udf.json;
/**
* 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 java.io.IOException;
import java.io.StringWriter;
import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
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.ConstantObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ListObjectInspector;
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.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector.PrimitiveCategory;
import org.apache.hadoop.hive.serde2.objectinspector.StructField;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.BinaryObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.BooleanObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.ByteObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.DoubleObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.FloatObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.IntObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.LongObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.ShortObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.StringObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.TimestampObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.WritableConstantBooleanObjectInspector;
import org.codehaus.jackson.JsonFactory;
import org.codehaus.jackson.JsonGenerationException;
import org.codehaus.jackson.JsonGenerator;
import org.joda.time.format.DateTimeFormatter;
import org.joda.time.format.ISODateTimeFormat;
/**
* Generate a JSON string from an arbitrary Hive structure.
* Use the struct() method to generate arbitrary JSON maps.
*
* <p>For example,
*
* to_json( struct("name":"Bob",
* "value",23.0,
* "color_list",array( "red", "yellow", "green" ),
* "inner_map", map( "a", 1, "b", 2, "c" 3) ) )
* = '{ "name":"Bob", "value":23.0,
* "color_list":[ "red", "yellow", "green" ],
* "inner_map":{ "a":1, "b":2, "c":3 } }'
*
*
*/
@Description(name="to_json",
value = "_FUNC_(struct, convert_to_camel_case) - Returns a JSON string from an arbitrary Hive structure."
)
public class ToJsonUDF extends GenericUDF {
private InspectorHandle inspHandle;
private Boolean convertFlag = Boolean.FALSE;
private JsonFactory jsonFactory;
private interface InspectorHandle {
abstract public void generateJson(JsonGenerator gen, Object obj) throws JsonGenerationException,IOException;
};
private class MapInspectorHandle implements InspectorHandle {
private MapObjectInspector mapInspector;
private StringObjectInspector keyObjectInspector;
private InspectorHandle valueInspector;
public MapInspectorHandle( MapObjectInspector mInsp) throws UDFArgumentException {
mapInspector = mInsp;
try {
keyObjectInspector = (StringObjectInspector) mInsp.getMapKeyObjectInspector();
} catch(ClassCastException castExc) {
throw new UDFArgumentException("Only Maps with strings as keys can be converted to valid JSON");
}
valueInspector = GenerateInspectorHandle( mInsp.getMapValueObjectInspector());
}
@Override
public void generateJson(JsonGenerator gen, Object obj) throws JsonGenerationException, IOException {
if( obj == null) {
gen.writeNull();
} else {
gen.writeStartObject();
Map map = mapInspector.getMap(obj);
Iterator<Map.Entry> iter = map.entrySet().iterator();
while( iter.hasNext()) {
Map.Entry entry = iter.next();
String keyJson = keyObjectInspector.getPrimitiveJavaObject(entry.getKey());
if( convertFlag) {
gen.writeFieldName( FromJsonUDF.ToCamelCase(keyJson));
} else {
gen.writeFieldName( keyJson);
}
valueInspector.generateJson( gen, entry.getValue() );
}
gen.writeEndObject();
}
}
}
private class StructInspectorHandle implements InspectorHandle {
private StructObjectInspector structInspector;
private List<String> fieldNames;
private List<InspectorHandle> fieldInspectorHandles;
public StructInspectorHandle(StructObjectInspector insp) throws UDFArgumentException {
structInspector = insp;
List<? extends StructField> fieldList = insp.getAllStructFieldRefs();
this.fieldNames = new ArrayList<String>();
this.fieldInspectorHandles = new ArrayList<InspectorHandle>();
for(StructField sf : fieldList) {
fieldNames.add( sf.getFieldName());
fieldInspectorHandles.add( GenerateInspectorHandle( sf.getFieldObjectInspector() ));
}
}
@Override
public void generateJson(JsonGenerator gen,Object obj) throws JsonGenerationException, IOException {
//// Interpret a struct as a map ...
if(obj == null) {
gen.writeNull();
} else {
gen.writeStartObject();
List structObjs = structInspector.getStructFieldsDataAsList(obj);
for(int i=0; i<fieldNames.size(); ++i) {
String fieldName = fieldNames.get(i);
if( convertFlag) {
gen.writeFieldName( FromJsonUDF.ToCamelCase(fieldName)) ;
} else {
gen.writeFieldName( fieldName);
}
fieldInspectorHandles.get(i).generateJson( gen, structObjs.get(i));
}
gen.writeEndObject();
}
}
}
private class ArrayInspectorHandle implements InspectorHandle {
private ListObjectInspector arrayInspector;
private InspectorHandle valueInspector;
public ArrayInspectorHandle( ListObjectInspector lInsp) throws UDFArgumentException {
arrayInspector = lInsp;
valueInspector = GenerateInspectorHandle( arrayInspector.getListElementObjectInspector());
}
@Override
public void generateJson(JsonGenerator gen, Object obj) throws JsonGenerationException, IOException {
if( obj == null) {
gen.writeNull();
} else {
gen.writeStartArray();
List list = arrayInspector.getList( obj);
for( Object listObj : list) {
valueInspector.generateJson(gen, listObj);
}
gen.writeEndArray();
}
}
}
private class StringInspectorHandle implements InspectorHandle {
private StringObjectInspector strInspector;
public StringInspectorHandle( StringObjectInspector insp) {
strInspector = insp;
}
@Override
public void generateJson(JsonGenerator gen, Object obj) throws JsonGenerationException, IOException {
if( obj == null) {
gen.writeNull();
} else {
String str = strInspector.getPrimitiveJavaObject(obj);
gen.writeString( str);
}
}
}
private class IntInspectorHandle implements InspectorHandle {
private IntObjectInspector intInspector;
public IntInspectorHandle( IntObjectInspector insp) {
intInspector = insp;
}
@Override
public void generateJson(JsonGenerator gen,Object obj) throws JsonGenerationException, IOException {
if( obj == null)
gen.writeNull();
else {
int num = intInspector.get(obj);
gen.writeNumber( num);
}
}
}
private class DoubleInspectorHandle implements InspectorHandle {
private DoubleObjectInspector dblInspector;
public DoubleInspectorHandle( DoubleObjectInspector insp) {
dblInspector = insp;
}
@Override
public void generateJson(JsonGenerator gen, Object obj) throws JsonGenerationException, IOException {
if( obj == null) {
gen.writeNull();
} else {
double num = dblInspector.get(obj);
gen.writeNumber( num);
}
}
}
private class LongInspectorHandle implements InspectorHandle {
private LongObjectInspector longInspector;
public LongInspectorHandle( LongObjectInspector insp) {
longInspector = insp;
}
@Override
public void generateJson(JsonGenerator gen,Object obj) throws JsonGenerationException, IOException {
if( obj == null) {
gen.writeNull();
} else {
long num = longInspector.get(obj);
gen.writeNumber( num);
}
}
}
private class ShortInspectorHandle implements InspectorHandle {
private ShortObjectInspector shortInspector;
public ShortInspectorHandle( ShortObjectInspector insp) {
shortInspector = insp;
}
@Override
public void generateJson(JsonGenerator gen,Object obj) throws JsonGenerationException, IOException {
if( obj == null) {
gen.writeNull();
} else {
short num = shortInspector.get(obj);
gen.writeNumber( num);
}
}
}
private class ByteInspectorHandle implements InspectorHandle {
private ByteObjectInspector byteInspector;
public ByteInspectorHandle( ByteObjectInspector insp) {
byteInspector = insp;
}
@Override
public void generateJson(JsonGenerator gen,Object obj) throws JsonGenerationException, IOException {
if( obj == null) {
gen.writeNull();
} else {
byte num = byteInspector.get(obj);
gen.writeNumber( num);
}
}
}
private class FloatInspectorHandle implements InspectorHandle {
private FloatObjectInspector floatInspector;
public FloatInspectorHandle( FloatObjectInspector insp) {
floatInspector = insp;
}
@Override
public void generateJson(JsonGenerator gen, Object obj) throws JsonGenerationException, IOException {
if( obj == null) {
gen.writeNull();
} else {
float num = floatInspector.get(obj);
gen.writeNumber(num);
}
}
}
private class BooleanInspectorHandle implements InspectorHandle {
private BooleanObjectInspector boolInspector;
public BooleanInspectorHandle( BooleanObjectInspector insp) {
boolInspector = insp;
}
@Override
public void generateJson(JsonGenerator gen, Object obj) throws JsonGenerationException, IOException {
if( obj == null) {
gen.writeNull();
} else {
boolean tf = boolInspector.get( obj);
gen.writeBoolean( tf);
}
}
}
private class BinaryInspectorHandle implements InspectorHandle {
private BinaryObjectInspector binaryInspector;
public BinaryInspectorHandle( BinaryObjectInspector insp) {
binaryInspector = insp;
}
@Override
public void generateJson(JsonGenerator gen, Object obj) throws JsonGenerationException, IOException {
if( obj == null) {
gen.writeNull();
} else {
byte[] bytes = binaryInspector.getPrimitiveJavaObject(obj);
gen.writeBinary(bytes);
}
}
}
private class TimestampInspectorHandle implements InspectorHandle {
private TimestampObjectInspector timestampInspector;
private DateTimeFormatter isoFormatter = ISODateTimeFormat.dateTimeNoMillis();
public TimestampInspectorHandle( TimestampObjectInspector insp) {
timestampInspector = insp;
}
@Override
public void generateJson(JsonGenerator gen, Object obj) throws JsonGenerationException, IOException {
if( obj == null) {
gen.writeNull();
} else {
Timestamp timestamp = timestampInspector.getPrimitiveJavaObject(obj);
String timeStr = isoFormatter.print( timestamp.getTime());
gen.writeString( timeStr);
}
}
}
private InspectorHandle GenerateInspectorHandle( ObjectInspector insp) throws UDFArgumentException {
Category cat = insp.getCategory();
if( cat == Category.MAP) {
return new MapInspectorHandle((MapObjectInspector) insp);
} else if( cat == Category.LIST ) {
return new ArrayInspectorHandle( (ListObjectInspector) insp);
} else if( cat == Category.STRUCT ) {
return new StructInspectorHandle( (StructObjectInspector) insp);
} else if( cat == Category.PRIMITIVE ) {
PrimitiveObjectInspector primInsp = (PrimitiveObjectInspector) insp;
PrimitiveCategory primCat = primInsp.getPrimitiveCategory();
if( primCat == PrimitiveCategory.STRING) {
return new StringInspectorHandle((StringObjectInspector) primInsp);
} else if( primCat == PrimitiveCategory.INT ) {
return new IntInspectorHandle((IntObjectInspector) primInsp);
} else if( primCat == PrimitiveCategory.LONG ) {
return new LongInspectorHandle((LongObjectInspector) primInsp);
} else if( primCat == PrimitiveCategory.SHORT ) {
return new ShortInspectorHandle((ShortObjectInspector) primInsp);
} else if( primCat == PrimitiveCategory.BOOLEAN) {
return new BooleanInspectorHandle((BooleanObjectInspector) primInsp);
} else if( primCat == PrimitiveCategory.FLOAT) {
return new FloatInspectorHandle((FloatObjectInspector) primInsp);
} else if( primCat == PrimitiveCategory.DOUBLE) {
return new DoubleInspectorHandle((DoubleObjectInspector) primInsp);
} else if( primCat == PrimitiveCategory.BYTE) {
return new ByteInspectorHandle((ByteObjectInspector)primInsp);
} else if( primCat == PrimitiveCategory.BINARY) {
return new BinaryInspectorHandle((BinaryObjectInspector)primInsp);
}
}
/// Dunno ...
throw new UDFArgumentException("Don't know how to handle object inspector " + insp);
}
@Override
public Object evaluate(DeferredObject[] args) throws HiveException {
try {
StringWriter writer = new StringWriter();
JsonGenerator gen = jsonFactory.createJsonGenerator( writer);
inspHandle.generateJson(gen, args[0].get() );
gen.close();
writer.close();
return writer.toString();
} catch( IOException io ) {
throw new HiveException(io);
}
}
@Override
public String getDisplayString(String[] args) {
return "to_json(" + args[0] + ")";
}
@Override
public ObjectInspector initialize(ObjectInspector[] args)
throws UDFArgumentException {
if(args.length != 1 && args.length != 2 ) {
throw new UDFArgumentException(" ToJson takes an object as an argument, and an optional to_camel_case flag");
}
ObjectInspector oi= args[0];
inspHandle = GenerateInspectorHandle( oi);
if(args.length == 2 ) {
ObjectInspector flagInsp = args[1];
if( flagInsp.getCategory() != Category.PRIMITIVE
|| ((PrimitiveObjectInspector)flagInsp).getPrimitiveCategory()
!= PrimitiveCategory.BOOLEAN
|| !(flagInsp instanceof ConstantObjectInspector )) {
throw new UDFArgumentException(" ToJson takes an object as an argument, and an optional to_camel_case flag");
}
WritableConstantBooleanObjectInspector constInsp= (WritableConstantBooleanObjectInspector) flagInsp;
convertFlag = constInsp.getWritableConstantValue().get();
}
jsonFactory = new JsonFactory();
return PrimitiveObjectInspectorFactory.javaStringObjectInspector;
}
}