How to avoid to broadcast a large lookup table in Spark - scala

Question

Can you help me to avoid broadcasting of a large lookup table? I have a table with measurements:
Measurement Value
x1 5.1
x2 8.9
x1 9.1
x3 4.4
x2 2.1
...
And a list of pairs:
P1 P2
x1 x2
x2 x3
...
The task is to get all values for both elements of every pair and put them into a magic function. That's how I solved it by broadcasting the large table with the measurements.
case class Measurement(measurement: String, value: Double)
case class Candidate(c1: String, c2: String)
val measurements = Seq(Measurement("x1", 5.1), Measurement("x2", 8.9),
Measurement("x1", 9.1), Measurement("x3", 4.4))
val candidates = Seq(Candidate("x1", "x2"), Candidate("x2", "x3"))
// create data frames
val dfm = sqc.createDataFrame(measurements)
val dfc = sqc.createDataFrame(candidates)
// broadcast lookup table
val lookup = sc.broadcast(dfm.rdd.map(r => (r(0), r(1))).collect())
// udf: run magic test with every candidate
val magic: ((String, String) => Double) = (c1: String, c2: String) => {
val lt = lookup.value
val c1v = lt.filter(_._1 == c1).map(_._2).map(_.asInstanceOf[Double])
val c2v = lt.filter(_._1 == c2).map(_._2).map(_.asInstanceOf[Double])
new Foo().magic(c1v, c2v)
}
val sq1 = udf(magic)
val dfks = dfc.withColumn("magic", sq1(col("c1"), col("c2")))
As you can guess I'm not pretty happy with the solution. For every pair I filter the lookup table twice, this isn't fast nor elegant. I'm using Spark 1.6.1.

Answer 1

An alternative would be to use RDD and join. Not sure what's better in term of performance though.
case class Measurement(measurement: String, value: Double)
case class Candidate(c1: String, c2: String)
val measurements = Seq(Measurement("x1", 5.1), Measurement("x2", 8.9),
Measurement("x1", 9.1), Measurement("x3", 4.4))
val candidates = Seq(Candidate("x1", "x2"), Candidate("x2", "x3"))
val rdm = sc.parallelize(measurements).map(r => (r.measurement, r.value)).groupByKey().cache()
val rdc = sc.parallelize(candidates).map(r => (r.c1, r.c2)).cache()
val firstColJoin = rdc.join(rdm).values
val secondColJoin = firstColJoin.join(rdm).values
secondColJoin.map { case (c1v, c2v) => new Foo().magic(c1v, c2v) }

Answer 2

Thank you for all comments. I read the comments, did some research and studied zero323 posts.
My current solution is using two joins and an UserDefinedAggregateFunction:
object GroupValues extends UserDefinedAggregateFunction {
def inputSchema = new StructType().add("x", DoubleType)
def bufferSchema = new StructType().add("buff", ArrayType(DoubleType))
def dataType = ArrayType(DoubleType)
def deterministic = true
def initialize(buffer: MutableAggregationBuffer) = {
buffer.update(0, ArrayBuffer.empty[Double])
}
def update(buffer: MutableAggregationBuffer, input: Row) = {
if (!input.isNullAt(0))
buffer.update(0, buffer.getSeq[Double](0) :+ input.getDouble(0))
}
def merge(buffer1: MutableAggregationBuffer, buffer2: Row) = {
buffer1.update(0, buffer1.getSeq[Double](0) ++ buffer2.getSeq[Double](0))
}
def evaluate(buffer: Row) = buffer.getSeq[Double](0)
}
// join data for candidate one
val j1 = dfc.join(dfm, dfc("c1") === dfm("measurement"))
// aggregate all c1 values to an array
val j1v = j1.groupBy(col("c1"), col("c2")).agg(GroupValues(col("value"))
.alias("c1-values"))
// join data for candidate two
val j2 = j1v.join(dfm, j1v("c2") === dfm("measurement"))
// aggregate all c2 values to an array
val j2v = j2.groupBy(col("c1"), col("c2"), col("c1-values"))
.agg(GroupValues(col("value")).alias("c2-values"))
Next step would be to use collect_list instead of UserDefinedAggregateFunction.