Calculate a mode for multiple columns - scala

Question

I would like to calculate a mode for multiple columns in the same time in Spark and use this calculated values to impute missings in a DataFrame. I found how to calculate e.g. a mean, but a mode is more complex I think.
Here is a mean calculation:
val multiple_mean = df.na.fill(df.columns.zip(
df.select(intVars.map(mean(_)): _*).first.toSeq
).toMap)
I am able to calculate a mode in brute force way:
var list = ArrayBuffer.empty[Float]
for(column <- df.columns){
list += df.select(column).groupBy(col(column)).count().orderBy(desc("count")).first.toSeq(0).asInstanceOf[Float]
}
val multiple_mode = df.na.fill(df.columns.zip(list.toSeq).toMap)
What way would be the best if we consider a performance?
Thank you for any help.

Answer 1

You could use UserDefinedAggregateFunction. The code below is tested in spark 1.6.2
First create a class which extends UserDefinedAggregateFunction.
import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types._
class ModeUDAF extends UserDefinedAggregateFunction{
override def dataType: DataType = StringType
override def inputSchema: StructType = new StructType().add("input", StringType)
override def deterministic: Boolean = true
override def bufferSchema: StructType = new StructType().add("mode", MapType(StringType, LongType))
override def initialize(buffer: MutableAggregationBuffer): Unit = {
buffer(0) = Map.empty[Any, Long]
}
override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
val buff0 = buffer.getMap[Any, Long](0)
val inp = input.get(0)
buffer(0) = buff0.updated(inp, buff0.getOrElse(inp, 0L) + 1L)
}
override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
val mp1 = buffer1.getMap[Any, Long](0)
val mp2 = buffer2.getMap[Any, Long](0)
buffer1(0) = mp1 ++ mp2.map { case (k, v) => k -> (v + mp1.getOrElse(k, 0L)) }
}
override def evaluate(buffer: Row): Any = {
lazy val st = buffer.getMap[Any, Long](0).toStream
val mode = st.foldLeft(st.head){case (e, s) => if (s._2 > e._2) s else e}
mode._1
}
}
Afterwords you could use it with your dataframe in the following manner.
val modeColumnList = List("some", "column", "names") // or df.columns.toList
val modeAgg = new ModeUDAF()
val aggCols = modeColumnList.map(c => modeAgg(df(c)))
val aggregatedModeDF = df.agg(aggCols.head, aggCols.tail: _*)
aggregatedModeDF.show()
Also you could use .collect on the final dataframe to collect the result in a scala data structure.
Note: The performance of this solution depends on the cardinality of the input column.