简体   繁体   中英

How to filter a dataframe in Spark scala with relational operators as variables?

I have a dataframe as below:

myDF:

+-----+
|value|
+-----+
|8    |
|8    |
|1    |
+-----+

The program reads from other computed dataframe and get the below two values:

val attr = 5
val opr = >

Now i need to filter myDF based on the values. So my result will be like below:

resultDF:
+-----+----------+
|value|result    |
+-----+----------+
|8    |GOOD      |
|8    |GOOD      |
|1    |BAD       |
+-----+----------+

Code I used:

val resultDF = myDF.withColumn("result", when(col("value") > attr, "GOOD").otherwise("BAD"))

Now, the attr and opr will change dynamically. Meaning the operator can be any of >, <, >=, <=, <> .

Based on the operator I receive my filter condition should change. Like I need to use the variable for the operator.

Can someone please advise ?

val resultDF = myDF.withColumn("result", when(col("value") opr attr, "GOOD").otherwise("BAD"))

Firstly, as @ Andrew said, it's bad idea to use dynamic sql without a big reason, because of undefined behavior, and difficulties in debugging. Assume you have joined values with operators dataframe, then you can use this code:

import spark.implicits._

val appData: DataFrame = Seq(
  ("1", ">"),
  ("1", ">"),
  ("3", "<="),
  ("4", "<>"),
  ("6", ">="),
  ("6", "==")
).toDF("value", "operator")

val attr = 5

def compare(value: String, operator: String, sample: Int): String = {
  val isValueCorrectForAttr: Boolean = operator match {
    case ">" => value.toInt > sample
    case "<" => value.toInt < sample
    case ">=" => value.toInt >= sample
    case "<=" => value.toInt <= sample
    case "==" => value.toInt == sample
    case "<>" => value.toInt != sample
    case _ => throw new IllegalArgumentException(s"Wrong operator: $operator")
  }
  if (isValueCorrectForAttr) "GOOD" else "BAD"
}

import org.apache.spark.sql.functions._
val dynamic_compare =  spark.udf.register("dynamic_compare", (v: String, op: String) => compare(v, op, attr))
appData.withColumn("result", dynamic_compare(col("value"), col("operator")))

if you don't have operator column, and just single operator, it can be more simple:

import spark.implicits._

val appData: DataFrame = Seq(
  "1",
  "1",
  "3",
  "4",
  "6",
  "6"
).toDF("value")

val attr = 5
val op = ">"

def compare(value: String, operator: String, sample: Int): String = {
  val isValueCorrectForAttr: Boolean = operator match {
    case ">" => value.toInt > sample
    case "<" => value.toInt < sample
    case ">=" => value.toInt >= sample
    case "<=" => value.toInt <= sample
    case "==" => value.toInt == sample
    case "<>" => value.toInt != sample
    case _ => throw new IllegalArgumentException(s"Wrong operator: $operator")
  }
  if (isValueCorrectForAttr) "GOOD" else "BAD"
}

import org.apache.spark.sql.functions._
val dynamic_compare =  spark.udf.register("dynamic_compare", (value: String) => compare(value, op, attr))
appData.withColumn("result", dynamic_compare(col("value")))

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM