scala - 在 Spark 数据帧中查找

Question

我正在使用 Spark 1.6，我想知道如何在数据帧中进行查找。

我有两个数据框员工和部门。

员工数据框

-------------------
Emp Id | Emp Name
------------------
1 | john
2 | David

部门数据框

--------------------
Dept Id | Dept Name | Emp Id
-----------------------------
1 | Admin | 1
2 | HR | 2

我想从员工表中查找 emp id 到部门表并获取部门名称。所以，结果集将是

Emp Id | Dept Name
-------------------
1 | Admin
2 | HR

我如何在 SPARK 中实现这个查找 UDF 功能。我不想在两个数据帧上都使用 JOIN。

Answer 1

正如评论中已经提到的，加入数据框是要走的路。

您可以使用查找，但我认为没有“分布式”解决方案，即您必须将查找表收集到驱动程序内存中。另请注意，此方法假定 EmpID 是唯一的:

import org.apache.spark.sql.functions._
import sqlContext.implicits._
import scala.collection.Map

val emp = Seq((1,"John"),(2,"David"))
val deps = Seq((1,"Admin",1),(2,"HR",2))

val empRdd = sc.parallelize(emp)
val depsDF = sc.parallelize(deps).toDF("DepID","Name","EmpID")


val lookupMap = empRdd.collectAsMap()
def lookup(lookupMap:Map[Int,String]) = udf((empID:Int) => lookupMap.get(empID))

val combinedDF = depsDF
  .withColumn("empNames",lookup(lookupMap)($"EmpID"))

我最初的想法是通过 empRdd到 UDF 并使用 lookup在 PairRDD 上定义的方法，但这当然不起作用，因为您不能在转换(即 UDF)中进行 Spark 操作(即 lookup )。

编辑:

如果您的 empDf 有多个列(例如 Name、Age)，您可以使用它

val empRdd = empDf.rdd.map{row =>
      (row.getInt(0),(row.getString(1),row.getInt(2)))}


    val lookupMap = empRdd.collectAsMap()
    def lookup(lookupMap:Map[Int,(String,Int)]) =
         udf((empID:Int) => lookupMap.lift(empID))

    depsDF
      .withColumn("lookup",lookup(lookupMap)($"EmpID"))
      .withColumn("empName",$"lookup._1")
      .withColumn("empAge",$"lookup._2")
      .drop($"lookup")
      .show()