CODEWITHSUNDEEP
scalaapache-sparkhadoopapache-spark-sql
Pavan_Obj
Pavan_Obj

Reputation: 1129

How can I optimize spark function to round a double value to 2 decimals?

Below is my Spark Function which is straight forward 

def doubleToRound(df:DataFrame,roundColsList:Array[String]): DataFrame ={
    var y:DataFrame = df
    for(colDF <- y.columns){
      if(roundColsList.contains(colDF)){
        y = y.withColumn(colDF,functions.round(y.col(colDF),2))
      }
    }

This is working as expected, by making the values of multiple columns for a given DF to round the decimal values to 2 positions. But I am looping through DataFrame y until the columns Array[Sting].length(). Any better way of doing the above?

Thank you All

Upvotes: 1

Views: 336

Answers (1)

Leo C
Leo C

Reputation: 22449

You can simply use select along with a map as shown in the following example:

import org.apache.spark.sql.functions._
import spark.implicits._

val df = Seq(
  ("a", 1.22, 2.333, 3.4444),
  ("b", 4.55, 5.666, 6.7777)
).toDF("id", "v1", "v2", "v3")

val roundCols = df.columns.filter(_.startsWith("v"))  // Or filter with other conditions
val otherCols = df.columns diff roundCols

df.select(otherCols.map(col) ++ roundCols.map(c => round(col(c), 2).as(c)): _*).show
// +---+----+----+----+
// | id|  v1|  v2|  v3|
// +---+----+----+----+
// |  a|1.22|2.33|3.44|
// |  b|4.55|5.67|6.78|
// +---+----+----+----+

Making it a method:

import org.apache.spark.sql.DataFrame

def doubleToRound(df: DataFrame, roundCols: Array[String]): DataFrame = {
  val otherCols = df.columns diff roundCols
  df.select(otherCols.map(col) ++ roundCols.map(c => round(col(c), 2).as(c)): _*)
}

Alternatively, use foldLeft and withColumn as follows:

def doubleToRound(df: DataFrame, roundCols: Array[String]): DataFrame =
  roundCols.foldLeft(df)((acc, c) => acc.withColumn(c, round(col(c), 2)))

Upvotes: 3

Related Questions