如何创建动态过滤条件并使用它来过滤 spark 数据帧上的行？

Question

There are two variables that I read from a file and use it to filter to a dataframe. Count of variables will vary on a daily basis so I need to build the filter condition dynamically so code change is not required.

clientid1=1,2,3,4
clientref1=10,20,30,40 
clientid2=1,2,3,4
clientref2=10,20,30,40 

I read the above file row by row and split by "=" and load it to Map[String,String]. In order to get the value, I pass the key "clientid1" and split the value by ","

val clientid1 :Array[String] = parameterValues("clientid1").split(",")

I used these variables like below to filter the DF and the below is working fine.

val finalDF = customerDF
              .filter((col("clientid").isin(clientid1: _*) && col("clientref").isin(clientref1: _*))
                      ||(col("clientid").isin(clientid2: _*) && col("clientref").isin(clientref2: _*)))

Now I need to build the below part dynamically,

(col("clientid").isin(clientid1: _*) && col("clientref").isin(clientref1: _*))
                          ||(col("clientid").isin(clientid2: _*) && col("clientref").isin(clientref2: _*))

if clientid3 and clientref3 are present in the input file then the filter condition should be

(col("clientid").isin(clientid1: _*) && col("clientref").isin(clientref1: _*))
||(col("clientid").isin(clientid2: _*) && col("clientref").isin(clientref2: _*))
||(col("clientid").isin(clientid3: _*) && col("clientref").isin(clientref3: _*))

I used a function like below to generate the condition but it is returning a String and the filter condition is not accepting it,

def generateFilterCond(count: Int, parameterValues: Map[String,String]): String ={
     var filterCond = ""
     for ( i <- 1 to count){
     val clientid :Array[String] = parameterValues("clientid$i").split(",")
     val clientref :Array[String] = parameterValues("clientref$i").split(",")

        filterCond = filterCond + (col("clientid").isin(clientid: _*) && col("clientref").isin(clientref: _*)) + " || " 
     }
        filterCond.substring(0, length-4).trim
}

Any suggestions on how to build this?

Answer 1

Indeed, filter take a spark Column as parameter but the method generateFilterCond returns a String . Thus the first step is to change signature to return a Column

However, we now have an issue because instead of initializing the filterCond variable in the first line of the function by an empty String , we need to initialize it with an empty Column , which does not exist in Spark.

But as we know that we will pass the column to a filter, and that the condition we will write is only a sequence of or , we can initialize this col to a condition that will always be false, such as lit(false) .

Thus, we can rewrite your method as follow:

def generateFilterCond(count: Int, parameterValues: Map[String,String]): Column = {
    var filterCond = lit(false)
    for (i <- 1 to count) {
      val clientid: Array[String] = parameterValues(s"clientid$i").split(",")
      val clientref: Array[String] = parameterValues(s"clientref$i").split(",")

      val addedCondition = col("clientid").isin(clientid: _*) && col("clientref").isin(clientref: _*)

      filterCond = filterCond || addedCondition
    }
    filterCond
  }

This method still have some issues. For instance, a NotSuchElementException exception will be raised if the parameterValues map doesn't contain a specific clientid/clientref between 1 and count . Moreover, it is not in a classic functional style, with a mutable accumulator and a loop.

So, assuming that we ignore all the problematic cases (no clientid/clientref couple or one of clientid or clientref doesn't exist but the other one exists), a final rewrite of this function could be:

  def generateFilterCond(count: Int, parameterValues: Map[String, String]): Column = (1 to count)
    .map(i => (parameterValues.get(s"clientid$i"), parameterValues.get(s"clientref$i")))
    .flatMap(elem => for {
      clientIds <- elem._1
      clientRefs <- elem._2
    } yield col("clientid").isin(clientIds.split(","): _*) && col("clientref").isin(clientRefs: _*)
    ).foldLeft(lit(false))((acc, c) => acc || c)

Be also aware that you can reach issues if the column expression resulting of your dynamic filter condition is too big, so you should keep your variables file small

Answer 2

To complete my previous answer , if your variables file contains a lot of lines, instead of trying to build a dynamic filter condition, you can read your variables file into a dataframe with spark and filter your input dataframe by joining it with the dataframe from your variables file

Filter dataframe without dynamic filter condition

In code below, your variable file is variables.txt , the dataframe on which you apply filter is input , and spark session is spark

import spark.implicits._

def extractInfo(line: String): (Int, String, Seq[Int]) = {
  val mainId = line.dropWhile(_.isLetter).takeWhile(_.isDigit).toInt
  val kind = line.takeWhile(_.isLetter)
  val values = line.dropWhile(s => s.isLetterOrDigit).tail.split(",").map(_.toInt)
  (mainId, kind, values)
}

def explodeIdRef(infos: Seq[(Int, String, Seq[Int])]): Seq[(Int, Int)] = for {
  clientId <- infos.find(_._2 == "clientid").map(_._3.toSeq).getOrElse(Nil)
  clientRef <- infos.find(_._2 == "clientref").map(_._3.toSeq).getOrElse(Nil)
} yield (clientId, clientRef)

val config = spark.read.text("variables.txt")
  .as[String]
  .map(extractInfo)
  .groupByKey(_._1)
  .flatMapGroups((_, iterator) => explodeIdRef(iterator.toSeq))
  .toDF("clientid", "clientref")

val filtered = input.join(config, Seq("clientid", "clientref"))

extractInfo will parse a line from variables.txt and extract useful columns. For instance, if the line is "clientid1=2,3" , extractInfo returns (1, "clientid", Seq(2, 3))

explodeIdRef will create all the couples clientid/clientref given a sequence of extracted info. For instance if you have Seq((1, "clientid", Seq(2, 3)), (1, "clientref", Seq(20, 30))) , explodeIdRef returns Seq((2, 20), (2, 30), (3, 20), (3, 30))

Test

if your variables.txt contains the following lines:

clientid1=1,2,3,4
clientref1=10,20,30,40
clientid2=5,6,7,8
clientref2=50,60,70,80

And you have the following input dataframe:

+---+--------+---------+
|id |clientid|clientref|
+---+--------+---------+
|1  |1       |10       |
|2  |2       |30       |
|3  |8       |10       |
|4  |8       |50       |
|5  |9       |90       |
+---+--------+---------+

The filtered dataframe will be:

+---+--------+---------+
|id |clientid|clientref|
+---+--------+---------+
|1  |1       |10       |
|2  |2       |30       |
|4  |8       |50       |
+---+--------+---------+