How to do aggregation on multiple columns at once in Spark

Question

I have a dataframe which has multiple columns. I want to group by one of the columns and aggregate other columns all the once. Let's say the table have 4 columns, cust_id, f1,f2,f3 and I want to group by cust_id and then get avg(f1), avg(f2) and avg(f3).The table will have many columns. Any hints?

The following code is good start but as I have many columns it may not be good idea to manually write them.

df.groupBy("cust_id").agg(sum("f1"), sum("f2"), sum("f3"))

Answer 1

Maybe you can try mapping a list with the colum names:

val groupCol = "cust_id"
val aggCols = (df.columns.toSet - groupCol).map(
  colName => avg(colName).as(colName + "_avg")
).toList

df.groupBy(groupCol).agg(aggCols.head, aggCols.tail: _*)

Alternatively, if needed, you can also match the schema and build the aggregations based on the type:

val aggCols = df.schema.collect {
  case StructField(colName, IntegerType, _, _) => avg(colName).as(colName + "_avg")
  case StructField(colName, StringType, _, _) => first(colName).as(colName + "_first")
}