How can I round a column in a Spark Dataset?

Viewed 6720

using Scala Spark, how can I use the typed Dataset API to round an aggregated column?

Also, how can I retain the type of a dataset through a groupby operation?

This is what I currently have:

case class MyRow(
  k1: String,
  k2: String,
  c1: Double,
  c2: Double
)

def groupTyped(ds: Dataset[MyRow]): Dataset[MyRow] = {
import org.apache.spark.sql.expressions.scalalang.typed._
ds.groupByKey(row => (row.k1, row.k2))
  .agg(
    avg(_.c1),
    avg(_.c2)
  )
  .map(r => MyRow(r._1._1, r._1._2, r._2, r._3))
}
  1. if I replace the avg(_.c1) with a round(avg(_.c1)) I get a type error. What's the correct way to round my values?
  2. the .map(...) line does not feel right -- is there a more elegant way to retain my Dataset's type?

Thanks!

2 Answers

While accepted answer works and is more general, in this case you can also use round. You just need to make column typed after rounding using .as[T] (also defining type to avg becomes necessary).

.agg(
  // Alternative ways to define a type to avg
  round(avg((r: MyRow) => r.c1)).as[Double],
  round(avg[MyRow](_.c2)).as[Double]
)
Related