CODEWITHSUNDEEP
scalaapache-sparkapache-spark-sql
Pavan_Obj
Pavan_Obj

Reputation: 1129

How can I optimize the spark function to replace nulls with zeroes?

Below is my Spark Function which handles nulls in a DataFrame column irrespective of its datatype.

  def nullsToZero(df:DataFrame,nullsToZeroColsList:Array[String]): DataFrame ={
    var y:DataFrame = df
    for(colDF <- y.columns){
      if(nullsToZeroColsList.contains(colDF)){
        y = y.withColumn(colDF,expr("case when "+colDF+" IS NULL THEN 0 ELSE "+colDF+" end"))
      }
    }
    return y
  }

    import spark.implicits._
    val personDF = Seq(
      ("miguel", Some(12),100,110,120), (null, Some(22),200,210,220), ("blu", None,300,310,320)
    ).toDF("name", "age","number1","number2","number3")
    println("Print Schema")
    personDF.printSchema()
    println("Show Original DF")
    personDF.show(false)
    val myColsList:Array[String] = Array("name","age","age")
    println("NULLS TO ZERO")
    println("Show NullsToZeroDF")
    val fixedDF = nullsToZero(personDF,myColsList)

In the above code I've a Integer type and a String type datatypes, both were handled by my funciton. But I suspect the below piece of code, in my function might affect the performance but not sure.

y = y.withColumn(colDF,expr("case when "+colDF+" IS NULL THEN 0 ELSE "+colDF+" end"))

Is there any more optimized way I can write this function, and what is the significance of doing .withColumn() and reassigning a DF again and again? Thank you in Advance.

Upvotes: 1

Views: 199

Answers (1)

Leo C
Leo C

Reputation: 22449

I would suggest assembling a valueMap for na.fill(valueMap) to fill null columns with specific values in accordance with the data types, as shown below:

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

val df = Seq(
  (Some(1), Some("a"), Some("x"), None),
  (None,    Some("b"), Some("y"), Some(20.0)),
  (Some(3), None,      Some("z"), Some(30.0))
).toDF("c1", "c2", "c3", "c4")

val nullColList = List("c1", "c2", "c4")

val valueMap = df.dtypes.filter(x => nullColList.contains(x._1)).
  collect{ case (c, t) => t match {
    case "StringType" => (c, "n/a")
    case "IntegerType" => (c, 0)
    case "DoubleType" => (c, Double.MinValue)
  } }.toMap
// valueMap: scala.collection.immutable.Map[String,Any] = 
//   Map(c1 -> 0, c2 -> n/a, c4 -> -1.7976931348623157E308)

df.na.fill(valueMap).show
// +---+---+---+--------------------+
// | c1| c2| c3|                  c4|
// +---+---+---+--------------------+
// |  1|  a|  x|-1.79769313486231...|
// |  0|  b|  y|                20.0|
// |  3|n/a|  z|                30.0|
// +---+---+---+--------------------+

Upvotes: 1

Related Questions