如何在多个条件下对Spark数据帧执行“查找”操作

时间:2016-03-01 09:40:13

标签: scala apache-spark dataframe apache-spark-sql lookup

我是Spark的新手(我的版本是1.6.0),现在我正在尝试解决下面给出的问题:

假设有两个源文件:

  • 第一个(简称A)是一个大的,包含名为A1,B1,C1和其他80列的列。里面有230K的记录。
  • 第二个(简称B)是一个小的查找表,其中包含名为A2,B2,C2和D2的列。里面有250条记录。

现在我们需要在A中插入一个新列,给出以下逻辑:

  • 首先在B中查找A1,B1和C1(对应列为A2,B2和C2),如果成功,则返回D2作为新添加列的值。如果没有找到......
  • 然后在B中查找A1,B1。如果成功,则返回D2。如果没有找到......
  • 设置默认值“NA”

我已经读过文件并将它们转换为数据框。对于第一种情况,我得到的结果是左外连接在一起。但是我在下一步找不到好的方法。

我目前的尝试是通过使用不太严格的条件连接A和B来构建新的数据框。但是我不知道如何从另一个更新当前数据帧。或者还有其他更直观,更有效的方法来解决整个问题吗?

感谢所有答案。

----------------------------- 20160309更新--------------- -----------------

终于接受了@mlk的回答。仍然非常感谢@ zero323对于他/她对UDF和加入的好评,Tungsten代码生成确实是我们现在面临的另一个问题。但是,由于我们需要为每次查找执行大量的查找和平均4个条件,因此前一个解决方案更合适......

最终解决方案在某种程度上看起来像下面的片段:

```
import sqlContext.implicits._
import com.github.marklister.collections.io._

case class TableType(A: String, B: String, C: String, D: String)
val tableBroadcast = sparkContext.broadcast(CsvParser(TableType).parseFile("..."))
val lkupD = udf {
  (aStr: String, bStr: String, cStr: String) =>
    tableBroadcast.value.find {
      case TableType(a, b, c, _) =>
        (a == aStr && b == bStr && c == cStr) ||
        (a == aStr && b == bStr)
    }.getOrElse(TableType("", "", "", "NA")).D
}
df = df.withColumn("NEW_COL", lkupD($"A", $"B", $"C"))
```

2 个答案:

答案 0 :(得分:4)

由于B很小,我认为最好的方法是广播变量和用户定义的功能。

// However you get the data...
case class BType( A2: Int, B2: Int, C2 : Int, D2 : String)
val B = Seq(BType(1,1,1,"B111"), BType(1,1,2, "B112"), BType(2,0,0, "B200"))

val A = sc.parallelize(Seq((1,1,1, "DATA"), (1,1,2, "DATA"), (2, 0, 0, "DATA"), (2, 0, 1, "NONE"), (3, 0, 0, "NONE"))).toDF("A1", "B1", "C1", "OTHER")


// Broadcast B so all nodes have a copy of it.
val Bbradcast = sc.broadcast(B)

// A user defined function to find the value for D2. This I'm sure could be improved by whacking it into maps. But this is a small example. 
val findD = udf {( a: Int, b : Int, c: Int) => Bbradcast.value.find(x => x.A2 == a && x.B2 == b && x.C2 == c).getOrElse(Bbradcast.value.find(x => x.A2 == a && x.B2 == b).getOrElse(BType(0,0,0,"NA"))).D2 }

// Use the UDF in a select
A.select($"A1", $"B1", $"C1", $"OTHER", findD($"A1", $"B1", $"C1").as("D")).show

答案 1 :(得分:2)

仅供参考没有UDF的解决方案:

val b1 = broadcast(b.toDF("A2_1", "B2_1", "C2_1", "D_1"))
val b2 = broadcast(b.toDF("A2_2", "B2_2", "C2_2", "D_2"))

// Match A, B and C
val expr1 = ($"A1" === $"A2_1") && ($"B1" === $"B2_1") && ($"C1" === $"C2_1")
// Match A and B mismatch C
val expr2 = ($"A1" === $"A2_2") && ($"B1" === $"B2_2") && ($"C1" !== $"C2_2")

val toDrop = b1.columns ++ b2.columns

toDrop.foldLeft(a
  .join(b1, expr1, "leftouter")
  .join(b2, expr2, "leftouter")
  // If there is match on A, B, C then D_1 should be not NULL
  // otherwise we fall-back to D_2 
  .withColumn("D", coalesce($"D_1", $"D_2")) 
)((df, c) => df.drop(c))

这假设每个类别中最多只有一个匹配(所有三列,或前两个)或输出中的重复行。

UDF vs JOIN

有多个因素需要考虑,这里没有简单的答案:

缺点

  • broadcast joins需要将数据两次传递给工作节点。至于现在broadcasted表没有缓存(SPARK-3863),并且在最近的将来不太可能发生变化(决议:稍后)。
  • 即使有完整匹配,
  • join操作也会应用两次。

赞成

  • joincoalesce对优化程序是透明的,而UDF则不是。
  • 直接使用SQL表达式操作可以从所有Tungsten优化中受益,包括代码生成,而UDF则不能。