Spark/Scala: use string variable in conditional expressions in DataFrame operations - scala

Question

Let me explain this with an example. Starting with the following dataframe
val df = Seq((1, "CS", 0, Array(0.1, 0.2, 0.4, 0.5)),
(4, "Ed", 0, Array(0.4, 0.8, 0.3, 0.6)),
(7, "CS", 0, Array(0.2, 0.5, 0.4, 0.7)),
(101, "CS", 1, Array(0.5, 0.7, 0.3, 0.8)),
(5, "CS", 1, Array(0.4, 0.2, 0.6, 0.9))).toDF("id", "dept", "test", "array")
df.show()
+---+----+----+--------------------+
| id|dept|test| array|
+---+----+----+--------------------+
| 1| CS| 0|[0.1, 0.2, 0.4, 0.5]|
| 4| Ed| 0|[0.4, 0.8, 0.3, 0.6]|
| 7| CS| 0|[0.2, 0.5, 0.4, 0.7]|
|101| CS| 1|[0.5, 0.7, 0.3, 0.8]|
| 5| CS| 1|[0.4, 0.2, 0.6, 0.9]|
+---+----+----+--------------------+
Considering the following two common operations as example (but do not have to be limited to them):
import org.apache.spark.sql.functions._ // for `when`
val dfFilter1 = df.where($"dept" === "CS")
val dfFilter3 = df.withColumn("category", when($"dept" === "CS" && $"id" === 101, 10).otherwise(0))
Now, I have a string variable colName = "dept". And it is required that $"dept" in the previous operation has to be replaced by colName in some form to achieve the same functionality. I managed to achieve the first one as following:
val dfFilter2 = df.where(s"${colName} = 'CS'")
But similar operation fails in the second case:
val dfFilter4 = df.withColumn("category", when(s"${colName} = 'CS'" && $"id" === 101, 10).otherwise(0))
Specifically it gives the following error:
Name: Unknown Error
Message: <console>:35: error: value && is not a member of String
val dfFilter4 = df.withColumn("category", when(s"${colName} = 'CS'" && $"id" === 101, 10).otherwise(0))
My understanding so far is that after I use s"${variable}" to deal with a variable, everything becomes pure string, and it is difficult to have logic operation involved.
So, my question are:
1. What is the best way to use such string variable as colName for operations similar as the two I listed above (I also do not like the solution I have for .where())?
2. Are there any general guidelines to use such string variable in more general operations other than the two examples here (I always felt that it is very case-specific when I deal with string related operations).

Answer 1

You can use expr function as
val dfFilter4 = df.withColumn("category", when(expr(s"${colName} = 'CS' and id = 101"), 10).otherwise(0))
Reason of the error
where function when defined with string query as following is working
val dfFilter2 = df.where(s"${colName} = 'CS'")
because there are supporting apis for both string and column
/**
* Filters rows using the given condition. This is an alias forfilter.
* {{{
* // The following are equivalent:
* peopleDs.filter($"age" > 15)
* peopleDs.where($"age" > 15)
* }}}
*
* #group typedrel
* #since 1.6.0
*/
def where(condition: Column): Dataset[T] = filter(condition)
and
/**
* Filters rows using the given SQL expression.
* {{{
* peopleDs.where("age > 15")
* }}}
*
* #group typedrel
* #since 1.6.0
*/
def where(conditionExpr: String): Dataset[T] = {
filter(Column(sparkSession.sessionState.sqlParser.parseExpression(conditionExpr)))
}
But there is only one api for when function supporting only column type
/**
* Evaluates a list of conditions and returns one of multiple possible result expressions.
* If otherwise is not defined at the end, null is returned for unmatched conditions.
*
* {{{
* // Example: encoding gender string column into integer.
*
* // Scala:
* people.select(when(people("gender") === "male", 0)
* .when(people("gender") === "female", 1)
* .otherwise(2))
*
* // Java:
* people.select(when(col("gender").equalTo("male"), 0)
* .when(col("gender").equalTo("female"), 1)
* .otherwise(2))
* }}}
*
* #group normal_funcs
* #since 1.4.0
*/
def when(condition: Column, value: Any): Column = withExpr {
CaseWhen(Seq((condition.expr, lit(value).expr)))
}
So you cannot use string sql query for when function
So, correct way of doing is as following
val dfFilter4 = df.withColumn("category", when(col(s"${colName}") === "CS" && $"id" === 101, 10).otherwise(0))
or in short as
val dfFilter4 = df.withColumn("category", when(col(colName) === "CS" && col("id") === 101, 10).otherwise(0))

Answer 2

What is the best way to use such string variable as colName for operations similar as the two I listed above
You can use col function from org.apache.spark.sql.functions
import org.apache.spark.sql.functions._
val colName = "dept"
For dfFilter2
val dfFilter2 = df.where(col(colName) === "CS")
For dfFilter4
val dfFilter4 = df.withColumn("category", when(col(colName) === "CS" && $"id" === 101, 10).otherwise(0))