CODEWITHSUNDEEP
scaladataframeapache-sparkapache-spark-sql
user3735871
user3735871

Reputation: 317

How to use Spark value in a column in if-else conditions - Scala

I have the following dataframe, where I hope to compare the value in a column to a string, and use it as a condition in if-else statement.

dataframe:

id   type       name
1    fruit     apple
2    toy       football
3   fruit      orange

what I am hoping to achieve is:

if(df("type") == "fruit"){
    //do something to the df
}else if ( df("type") == "toy"){
    //something else to the df
}

I tried to use val type= df.select("type").collectAsList().getString(0) but that's not correct. Could someone please help? Many thanks. I don't think it's a duplicate to this question as I do not want to add a new column. Spark: Add column to dataframe conditionally and I do not wish to use withColumn

Upvotes: 1

Views: 3052

Answers (1)

Sathvik
Sathvik

Reputation: 25

Implicit class for Dataframe should do the trick, code below (Ignore my crazy imports😀)

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

implicit class typeTransform(df: DataFrame) {
    def tranformByType(typeCol: Column) = {
      if(df.filter(typeCol === lit("fruit")).count > 0) df // do something to df 
      //........ more if Statements
      else df// do something with df 
    }
}

Usage can be something like this

val someDf = Seq(("fruit", "1"), ("toy", "2")).toDF("type", "id").tranformByType(col("type"))

Upvotes: 1

Related Questions