title: Spark: Efficient mass lookup in pair RDD's
body: <p>In Apache Spark I have two RDD's. The first <code>data : RDD[(K,V)]</code> containing data in key-value form. The second <code>pairs : RDD[(K,K)]</code> contains a set of interesting key-pairs of this data.</p>
  
  <p><strong>How can I efficiently construct an RDD <code>pairsWithData : RDD[((K,K)),(V,V))]</code></strong>, such that it contains all the elements from <code>pairs</code> as the key-tuple and their corresponding values (from <code>data</code>) as the value-tuple?</p>
  
  <p><strong>Some properties of the data:</strong></p>
  
  <ul>
  <li>The keys in <code>data</code> are unique</li>
  <li>All entries in <code>pairs</code> are unique</li>
  <li>For all pairs <code>(k1,k2)</code> in <code>pairs</code> it is guaranteed that <code>k1 &lt;= k2</code></li>
  <li>The size of 'pairs' is only a constant the size of data <code>|pairs| = O(|data|)</code></li>
  <li>Current data sizes (expected to grow): <code>|data| ~ 10^8, |pairs| ~ 10^10</code></li>
  </ul>
  
  <h2>Current attempts</h2>
  
  <p>Here is some example code in Scala:</p>
  
  <pre><code>import org.apache.spark.rdd.RDD
  import org.apache.spark.SparkContext._
  
  // This kind of show the idea, but fails at runtime.
  def massPairLookup1(keyPairs : RDD[(Int, Int)], data : RDD[(Int, String)]) = {
    keyPairs map {case (k1,k2) =&gt;
      val v1 : String = data lookup k1 head;
      val v2 : String = data lookup k2 head;
      ((k1, k2), (v1,v2))
    }
  }
  
  // Works but is O(|data|^2)
  def massPairLookup2(keyPairs : RDD[(Int, Int)], data : RDD[(Int, String)]) = {
    // Construct all possible pairs of values
    val cartesianData = data cartesian data map {case((k1,v1),(k2,v2)) =&gt; ((k1,k2),(v1,v2))}
    // Select only the values who's keys are in keyPairs
    keyPairs map {(_,0)} join cartesianData mapValues {_._2}
  }
  
  // Example function that find pairs of keys
  // Runs in O(|data|) in real life, but cannot maintain the values
  def relevantPairs(data : RDD[(Int, String)]) = {
    val keys = data map (_._1)
    keys cartesian keys filter {case (x,y) =&gt; x*y == 12 &amp;&amp; x &lt; y}
  }
  
  // Example run
  val data = sc parallelize(1 to 12) map (x =&gt; (x, "Number " + x))
  val pairs = relevantPairs(data)
  val pairsWithData = massPairLookup2(pairs, data) 
  
  
  // Print: 
  // ((1,12),(Number1,Number12))
  // ((2,6),(Number2,Number6))
  // ((3,4),(Number3,Number4))
  pairsWithData.foreach(println)
  </code></pre>
  
  <p><strong>Attempt 1</strong></p>
  
  <p>First I tried just using the <code>lookup</code> function on <code>data</code>, but that throws an runtime error when executed. It seems like <code>self</code> is null in the <code>PairRDDFunctions</code> trait.</p>
  
  <p>In addition I am not sure about the performance of <code>lookup</code>. <a href="https://spark.apache.org/docs/latest/api/scala/#org.apache.spark.rdd.PairRDDFunctions" rel="nofollow">The documentation</a> says <em>This operation is done efficiently if the RDD has a known partitioner by only searching the partition that the key maps to.</em> This sounds like <code>n</code> lookups takes O(n*|partition|) time at best, which I suspect could be optimized.</p>
  
  <p><strong>Attempt 2</strong></p>
  
  <p>This attempt works, but I create <code>|data|^2</code> pairs which will kill performance. I do not expect Spark to be able to optimize that away.</p>