topByKey.scala

  def topByKey(key: String, orderBy: String, n: Int): DataFrame = {
    val keyIndex = df.schema.fieldIndex(key)
    val orderByIndex = df.schema.fieldIndex(orderBy)
    val ord = df.schema.fields(orderByIndex).dataType match {
      case o: StringType => Ordering.by[Row, String](_.getString(orderByIndex))
      case o: IntegerType => Ordering.by[Row, Int](_.getInt(orderByIndex))
      case o: LongType => Ordering.by[Row, Long](_.getLong(orderByIndex))
      case o: FloatType => Ordering.by[Row, Float](_.getFloat(orderByIndex))
      case o: DoubleType => Ordering.by[Row, Double](_.getDouble(orderByIndex))
      case _ => throw new IllegalArgumentException
    }
    val rdd = df
      .map { row =>
        (row.get(keyIndex), row)
      }
      .topByKey(n)(ord)
      .flatMap(_._2)
    df.sqlContext.createDataFrame(rdd, df.schema)
  }