从Spark Dataframe构建2D查找表

时间:2017-09-11 23:20:54

标签: scala apache-spark dataframe broadcast lookup-tables

我想将较小的数据帧转换为广播查找表,以便在另一个较大的数据帧的UDF中使用。这个较小的数据框( myLookupDf )可能如下所示:

+---+---+---+---+
| x | 90|100|101|
+---+---+---+---+
| 90|  1|  0|  0|
|100|  0|  1|  1|
|101|  0|  1|  1|
+---+---+---+---+

我想使用第一列作为第一个键,比如x1,第一行作为第二个键。 x1和x2具有相同的元素。理想情况下,查找表( myLookupMap )将是Scala Map(或类似),其工作方式如下:

myLookupMap(90)(90) returns 1
myLookupMap(90)(101) returns 0
myLookupMap(100)(90) returns 0
myLookupMap(101)(100) return 1
etc.

到目前为止,我设法:

val myLookupMap = myLookupDf.collect().map(r => Map(myLookupDf.columns.zip(r.toSeq):_*))
myLookupMap: Array[scala.collection.Map[String,Any]] = Array(Map(x -> 90, 90 -> 1, 100 -> 0, 101 -> 0), Map(x -> 100, 90 -> 0, 100 -> 1, 101 -> 1), Map(x -> 101, 90 -> 0, 100 -> 1, 101 -> 1))

这是一个Map数组,并不完全是必需的。任何建议都非常感谢。

1 个答案:

答案 0 :(得分:1)

collect()始终创建rdd,相当于Array。您必须找到将arrays收集为maps的方法。

鉴于dataframe

scala> myLookupDf.show(false)
+---+---+---+---+
|x  |90 |100|101|
+---+---+---+---+
|90 |1  |0  |0  |
|100|0  |1  |1  |
|101|0  |1  |1  |
+---+---+---+---+

您需要的只是x以外的标题名称,因此您可以执行以下操作

scala>     val header = myLookupDf.schema.fieldNames.tail
header: Array[String] = Array(90, 100, 101)

我只是修改您的map函数以获取Map作为结果

scala>     val myLookupMap = myLookupDf.rdd.map(r => {
     |       val row = r.toSeq
     |       (row.head, Map(header.zip(row.tail):_*))
     |     }).collectAsMap()
myLookupMap: scala.collection.Map[Any,scala.collection.immutable.Map[String,Any]] = Map(101 -> Map(90 -> 0, 100 -> 1, 101 -> 1), 100 -> Map(90 -> 0, 100 -> 1, 101 -> 1), 90 -> Map(90 -> 1, 100 -> 0, 101 -> 0))

你应该看到你得到了理想的结果。

scala> myLookupMap(90)(90.toString)
res1: Any = 1

scala> myLookupMap(90)(101.toString)
res2: Any = 0

scala> myLookupMap(100)(90.toString)
res3: Any = 0

scala> myLookupMap(101)(100.toString)
res4: Any = 1

现在,您可以将myLookupMap传递给udf函数