import java.util.ArrayList;

import java.util.HashMap;

import java.util.List;

import java.util.Map;

import org.apache.spark.sql.Row;

import org.apache.spark.sql.expressions.MutableAggregationBuffer;

import org.apache.spark.sql.expressions.UserDefinedAggregateFunction;

import org.apache.spark.sql.types.DataType;

import org.apache.spark.sql.types.DataTypes;

import org.apache.spark.sql.types.StructField;

import org.apache.spark.sql.types.StructType;

public class SparkUDAF extends UserDefinedAggregateFunction

{

private StructType inputSchema;

private StructType bufferSchema;

private DataType returnDataType =

DataTypes.createMapType(DataTypes.StringType, DataTypes.StringType);

MutableAggregationBuffer mutableBuffer;

public SparkUDAF()

{

//inputSchema : This UDAF can accept 2 inputs which are of type Integer

List<StructField> inputFields = new ArrayList<StructField>();

StructField inputStructField1 = DataTypes.createStructField(“femaleCount”,DataTypes.IntegerType, true);

inputFields.add(inputStructField1);

StructField inputStructField2 = DataTypes.createStructField(“maleCount”,DataTypes.IntegerType, true);

inputFields.add(inputStructField2);

inputSchema = DataTypes.createStructType(inputFields);

//BufferSchema : This UDAF can hold calculated data in below mentioned buffers

List<StructField> bufferFields = new ArrayList<StructField>();

StructField bufferStructField1 = DataTypes.createStructField(“totalCount”,DataTypes.IntegerType, true);

bufferFields.add(bufferStructField1);

StructField bufferStructField2 = DataTypes.createStructField(“femaleCount”,DataTypes.IntegerType, true);

bufferFields.add(bufferStructField2);

StructField bufferStructField3 = DataTypes.createStructField(“maleCount”,DataTypes.IntegerType, true);

bufferFields.add(bufferStructField3);

StructField bufferStructField4 = DataTypes.createStructField(“outputMap”,DataTypes.createMapType(DataTypes.StringType, DataTypes.StringType), true);

bufferFields.add(bufferStructField4);

bufferSchema = DataTypes.createStructType(bufferFields);

}

/**

* This method determines which bufferSchema will be used

*/

@Override

public StructType bufferSchema() {

return bufferSchema;

}

/**

* This method determines the return type of this UDAF

*/

@Override

public DataType dataType() {

return returnDataType;

}

/**

* Returns true iff this function is deterministic, i.e. given the same input, always return the same output.

*/

@Override

public boolean deterministic() {

return true;

}

/**

* This method will re-initialize the variables to 0 on change of city name

*/

@Override

public void initialize(MutableAggregationBuffer buffer) {

buffer.update(, );

buffer.update(, );

buffer.update(, );

mutableBuffer = buffer;

}

/**

* This method is used to increment the count for each city

*/

@Override

public void update(MutableAggregationBuffer buffer, Row input) {

buffer.update(, buffer.getInt() + input.getInt() + input.getInt());

buffer.update(, input.getInt());

buffer.update(, input.getInt());

}

/**

* This method will be used to merge data of two buffers

*/

@Override

public void merge(MutableAggregationBuffer buffer, Row input) {

buffer.update(, buffer.getInt() + input.getInt());

buffer.update(, buffer.getInt() + input.getInt());

buffer.update(, buffer.getInt() + input.getInt());

}

/**

* This method calculates the final value by referring the aggregation buffer

*/

@Override

public Object evaluate(Row buffer) {

//In this method we are preparing a final map that will be returned as output

Map<String,String> op = new HashMap<String,String>();

op.put(“Total”, “” + mutableBuffer.getInt());

op.put(“dominant”, “Male”);

if(buffer.getInt() > mutableBuffer.getInt())

{

op.put(“dominant”, “Female”);

}

mutableBuffer.update(,op);

return buffer.getMap();

}

/**

* This method will determine the input schema of this UDAF

*/

@Override

public StructType inputSchema() {

return inputSchema;

}

}

Now lets see how we can access this UDAF using our spark code

import java.util.ArrayList;

import java.util.HashMap;

import java.util.List;

import java.util.Map;

import java.util.StringTokenizer;

import org.apache.spark.SparkConf;

import org.apache.spark.api.java.JavaRDD;

import org.apache.spark.api.java.JavaSparkContext;

import org.apache.spark.api.java.function.Function;

import org.apache.spark.sql.DataFrame;

import org.apache.spark.sql.Row;

import org.apache.spark.sql.RowFactory;

import org.apache.spark.sql.SQLContext;

import org.apache.spark.sql.hive.HiveContext;

import org.apache.spark.sql.types.DataTypes;

import org.apache.spark.sql.types.StructField;

import org.apache.spark.sql.types.StructType;

public class TestDemo {

public static void main (String args[])

{

//Set up sparkContext and SQLContext

SparkConf conf = new SparkConf().setAppName(“udaf”).setMaster(“local”);

JavaSparkContext sc = new JavaSparkContext(conf);

SQLContext sqlContext = new org.apache.spark.sql.SQLContext(sc);

//create Row RDD

JavaRDD<String> citiesRdd = sc.textFile(“cities.txt”);

JavaRDD<Row> rowRdd = citiesRdd.map(new Function<String, Row>() {

public Row call(String line) throws Exception {

StringTokenizer st = new StringTokenizer(line,” “);

return RowFactory.create(st.nextToken().trim(),Integer.parseInt(st.nextToken().trim()),Integer.parseInt(st.nextToken().trim()));

}

});

//Create Struct Type

List<StructField> inputFields = new ArrayList<StructField>();

StructField inputStructField = DataTypes.createStructField(“city”,DataTypes.StringType, true);

inputFields.add(inputStructField);

StructField inputStructField2 = DataTypes.createStructField(“Female”,DataTypes.IntegerType, true);

inputFields.add(inputStructField2);

StructField inputStructField3 = DataTypes.createStructField(“Male”,DataTypes.IntegerType, true);

inputFields.add(inputStructField3);

StructType inputSchema = DataTypes.createStructType(inputFields);

//Create Data Frame

DataFrame df = sqlContext.createDataFrame(rowRdd, inputSchema);

//Register our Spark UDAF

SparkUDAF sparkUDAF = new SparkUDAF();

sqlContext.udf().register(“uf”,sparkUDAF);

//Register dataframe as table

df.registerTempTable(“cities”);

//Run query

sqlContext.sql(“SELECT city , count[‘dominant’] as Dominant, count[‘Total’] as Total from(select city, uf(Female,Male) as count from cities group by (city)) temp”).show(false);

}

}