在Spark数据框中提取列的值

Question

I have a requirement , where I need to filter out rows from spark dataframe where value of a certain column (say "price") needs to be matched with values present in a scala map.The key of scala map is value of another column (say "id"). My dataframe contains two columns : id and price. I need to filter out all the columns where price does not match the price mentioned in scala map.

My code resembles this:

object obj1{
  // This method returns value price for items as per their id
  getPrice(id:String):String {
   //lookup in a map and return the price
  }
}

object Main{    
  val validIds = Seq[String]("1","2","3","4")
  val filteredDf = baseDataframe.where(baseDataframe("id").in(validIDs.map(lit(_)): _*) &&
    baseDataframe("price") === (obj1.getPrice(baseDataframe("id").toString()))) 

  // But this line send string "id" to obj1.getPrice() function
  // rather than value of id column
  }
}

I am not able to pass value of id columns to function obj1.getPrice(). Any suggestion how to achieve this?

Thanks,

Answer 1

You can write a udf to do this:

val checkPrice(id: String, price: String) = validIds.exists(_ == id) && obj1.getPrice(id) == price
val checkPriceUdf = udf(checkPrice)

baseDataFrame.where(checkPriceUdf($"id", $"price"))

Or another solution is convert the Map of id -> price to a data frame, and then do an inner join with baseDataFrame on the id and price columns.