public class GenericUDAFTopNRow extends AbstractGenericUDAFResolver {
@Override
public GenericUDAFEvaluator getEvaluator(TypeInfo[] parameters)
throws SemanticException {
if (parameters.length < 2) {
throw new UDFArgumentTypeException(parameters.length - 1,"At least two argument is expected.");
}
if(!(TypeInfoUtils.getStandardWritableObjectInspectorFromTypeInfo(parameters[0]) instanceof WritableIntObjectInspector)){
throw new UDFArgumentTypeException(0,"The first argument must be integer,"+TypeInfoUtils.getStandardWritableObjectInspectorFromTypeInfo(parameters[0]).getClass());
}
if (!ObjectInspectorUtils.compareSupported(TypeInfoUtils.getStandardWritableObjectInspectorFromTypeInfo(parameters[1]))) {
throw new UDFArgumentTypeException(1,"Cannot support comparison of map<> type or complex type containing map<>.");
}
return new TopNEvaluator();
}
static class TopNBuffer implements AggregationBuffer {
List<Object[]> container;
}
public static class TopNEvaluator extends GenericUDAFEvaluator {
int size;
String[] fieldNM;
ObjectInspector[] fieldOI;
ObjectInspector[] originalOI;
StandardListObjectInspector partialOI;
StandardStructObjectInspector partialElemOI;
@Override
public ObjectInspector init(Mode m, ObjectInspector[] parameters)
throws HiveException {
super.init(m, parameters);
if (m == Mode.PARTIAL1 || m == Mode.COMPLETE) {
this.originalOI = new ObjectInspector[parameters.length];
System.arraycopy(parameters, 0, this.originalOI, 0, parameters.length);
this.size = parameters.length-1;
this.fieldNM = new String[this.size];
this.fieldOI = new ObjectInspector[this.size];
for (int i = 0; i < this.size; i++) {
this.fieldNM[i] = "f" + i;
this.fieldOI[i] = ObjectInspectorUtils.getStandardObjectInspector(parameters[i+1]);
}
return ObjectInspectorFactory.getStandardListObjectInspector(ObjectInspectorFactory.getStandardStructObjectInspector(Arrays.asList(this.fieldNM),Arrays.asList(this.fieldOI)));
} else if (m == Mode.PARTIAL2 || m == Mode.FINAL) {
this.partialOI = (StandardListObjectInspector) parameters[0];
this.partialElemOI=(StandardStructObjectInspector) this.partialOI.getListElementObjectInspector();
List<? extends StructField> structFieldRefs = this.partialElemOI.getAllStructFieldRefs();
this.size = structFieldRefs.size();
this.fieldNM = new String[this.size];
this.fieldOI = new ObjectInspector[this.size];
for (int i = 0; i < this.size; i++) {
StructField sf = structFieldRefs.get(i);
this.fieldNM[i] = sf.getFieldName();
this.fieldOI[i] = sf.getFieldObjectInspector();
}
return ObjectInspectorUtils.getStandardObjectInspector(this.partialOI);
}
return null;
}
@Override
public AggregationBuffer getNewAggregationBuffer() throws HiveException {
TopNBuffer buffer = new TopNBuffer();
reset(buffer);
return buffer;
}
@Override
public void reset(AggregationBuffer agg) throws HiveException {
TopNBuffer buffer = (TopNBuffer) agg;
buffer.container = new LinkedList<Object[]>();
}
@Override
public void iterate(AggregationBuffer agg, Object[] parameters)
throws HiveException {
/*如果查询结果为空,不作处理*/
if(isEmptySet(agg,parameters)){
return;
}
TopNBuffer buffer = (TopNBuffer) agg;
int n = ((WritableIntObjectInspector)this.originalOI[0]).get(parameters[0]);
int s = buffer.container.size();
if(s < n){
Object[] elemVal = new Object[this.size];
for (int j = 0; j < this.size; j++) {
elemVal[j] = ObjectInspectorUtils.copyToStandardObject(parameters[j+1],this.originalOI[j+1]);
}
buffer.container.add(elemVal);
/*make sure the size should be n*/
while(buffer.container.size() < n){
buffer.container.add(new Object[this.size]);
}
}else{
for(int i = 0;i < s;i++){
if (ObjectInspectorUtils.compare(buffer.container.get(i)[0],this.fieldOI[0], parameters[1], this.originalOI[1]) < 0) {
Object[] elemVal = new Object[this.size];
for(int j=0;j<this.size;j++){
elemVal[j] = ObjectInspectorUtils.copyToStandardObject(parameters[j+1],this.originalOI[j+1]);
}
buffer.container.add(i, elemVal);
break;
}
}
/*make sure the size should be n*/
while(buffer.container.size() > n){
buffer.container.remove(n);
}
}
}
@Override
public Object terminatePartial(AggregationBuffer agg)
throws HiveException {
TopNBuffer buffer = (TopNBuffer) agg;
return buffer.container.isEmpty()?null:buffer.container;
}
@Override
public void merge(AggregationBuffer agg, Object partial) {
/*如果查询结果为空,不作处理*/
if(isEmptySet(agg,partial)){
return;
}
TopNBuffer buffer = (TopNBuffer) agg;
List<?> listVal = this.partialOI.getList(partial);
final int cn = Math.max(buffer.container.size(), listVal.size());
List<Object[]> values = new LinkedList<Object[]>();
for(Object elemObj:listVal){
List<Object> elemVal=this.partialElemOI.getStructFieldsDataAsList(elemObj);
Object[] value=new Object[this.size];
for(int i=0,n=elemVal.size();i<n;i++){
value[i]=ObjectInspectorUtils.copyToStandardObject(elemVal.get(i), this.fieldOI[i]);
}
values.add(value);
}
buffer.container=mergeSortNotNull(buffer.container, values);
while(buffer.container.size()<cn){
buffer.container.add(new Object[this.size]);
}
while(buffer.container.size() > cn){
buffer.container.remove(cn);
}
}
@Override
public Object terminate(AggregationBuffer agg) throws HiveException {
TopNBuffer buffer = (TopNBuffer) agg;
return buffer.container.isEmpty()?null:buffer.container;
}
private List<Object[]> mergeSortNotNull(List<Object[]> list1, List<Object[]> list2){
List<Object[]> result=new LinkedList<Object[]>();
int i1=0, i2=0, n1=list1.size(), n2=list2.size();
while(i1<n1 && i2<n2){
if(list1.get(i1)[0]==null){
i1++;
continue;
}
if(list2.get(i2)[0]==null){
i2++;
continue;
}
int cp = ObjectInspectorUtils.compare(list1.get(i1)[0],this.fieldOI[0], list2.get(i2)[0], this.fieldOI[0]);
if(cp > 0){
result.add(list1.get(i1));
i1++;
}else if(cp<0){
result.add(list2.get(i2));
i2++;
}else{
result.add(list1.get(i1));
i1++;
i2++;
}
}
while(i1<n1){
if(list1.get(i1)[0]==null){
i1++;
continue;
}
result.add(list1.get(i1));
i1++;
}
while(i2<n2){
if(list2.get(i2)[0]==null){
i2++;
continue;
}
result.add(list2.get(i2));
i2++;
}
return result;
}
private boolean isEmptySet(AggregationBuffer agg, Object[] parameters){
if(agg==null || parameters==null){
return true;
}else{
for(int i=0; i<parameters.length; i++){
if(parameters[i]!=null){
return false;
}
}
return true;
}
}
private boolean isEmptySet(AggregationBuffer agg, Object parameter){
return (agg==null) || (parameter==null);
}
}
}