在Spark數據框中提取列的值

Question

我有一個要求，在這里我需要從Spark數據框中過濾出某些列（例如“價格”）的值需要與scala映射中存在的值匹配的行。scala映射的關鍵是另一列的值（說“ id”）。 我的數據框包含兩列：id和price。 我需要過濾出所有價格與scala地圖中提到的價格不匹配的列。

我的代碼類似於：

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
  }
}

我無法將id列的值傳遞給函數obj1.getPrice（）。 任何建議如何實現這一目標？

謝謝，

Answer 1

您可以編寫udf來做到這一點：

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

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

另一個解決方案是將id-> price的Map轉換為數據框，然后使用id和price列上的baseDataFrame進行內部baseDataFrame 。