title: Problems with recursive algorithms and Spark DataFrames
body: <p>If you want to do anything interesting with Graphs -- whether with <code>GraphX</code> or the new <code>GraphFrames</code> -- you end up doing recursive algorithms. The problem I am having is that when using <code>DataFrames</code> each iteration of the algorithm takes longer and longer, and more execution stages get kicked off per iteration. I'm more of a functional Spark user -- I can make things happen but don't have a full grasp of what happens under the hood. But my guess is that the lineage chain keeps extending, and without breaking the lineage chain, the earlier iterations get re-calculated with each step. So iteration 1 does iteration 1; iteration 2 does iteration 1 again, then iteration 2; iteration 3 has to do 1, then 2 again, etc.</p>
  
  <p>So my first question: Is that really what's happening, more or less?</p>
  
  <p>To test it, I've been playing with <code>RDD.checkpoint</code>. It seems to help, but I can't really prove it. So that's my second question -- tell me if the way I use <code>checkpoint</code> is helping or not.</p>
  
  <p>Lastly, it would be great to hear about other possible solutions. Maybe Spark isn't even the right answer. I'm open to anything.</p>
  
  <p>To test all of this, I have been using a simple algorithm to polyfill vertex attributes -- a sort of attribute inheritance. Given a graph like this:</p>
  
  <pre><code>val nodes = Seq(
      (1L, Option(1L), Option(1L)),
      (2L, None, Option(2L)),
      (3L, Option(2L), None),
      (4L, None, None)
  ).toDF("id","inputType","recurrence")
  
  val edges = Seq(
      (1L, 2L, "parent"),
      (2L, 4L, "parent"),
      (1L, 3L, "parent")
  ).toDF("src","dst","type")
  </code></pre>
  
  <p>Once I polyfill missing attributes in the vertices, I should get something like this:</p>
  
  <pre><code>+---+---------+----------+
  | id|inputType|recurrence|
  +---+---------+----------+
  |  1|        1|         1|
  |  2|        1|         2|
  |  3|        2|         1|
  |  4|        1|         2|
  +---+---------+----------+
  </code></pre>
  
  <p>Vertex <code>1L</code> is the parent, and the other vertices have inherited missing attributes from the parent, following the chain upwards if needed.</p>
  
  <p>The algorithm actually isn't very complicated -- I'm going to use my own cobbled-together DataFrame / Graph algorithms instead of <code>GraphFrames</code>.</p>
  
  <p>First, I'll define a function to create edge triplets out of nodes and edges:</p>
  
  <pre><code>import org.apache.spark.sql.DataFrame
  def triplets(vertices: DataFrame, edges: DataFrame) : DataFrame = {
    edges.toDF(edges.columns.map(c =&gt; "edge_" + c):_*)
      .join(vertices.toDF(vertices.columns.map(c =&gt; "src_" + c):_*), col("edge_src") === col("src_id"))
      .join(vertices.toDF(vertices.columns.map(c =&gt; "dst_" + c):_*), col("edge_dst") === col("dst_id"))
  }
  </code></pre>
  
  <p>Based on the data above, <code>triplets(nodes,edges)</code> shows:</p>
  
  <pre><code>+--------+--------+---------+------+-------------+--------------+------+-------------+--------------+
  |edge_src|edge_dst|edge_type|src_id|src_inputType|src_recurrence|dst_id|dst_inputType|dst_recurrence|
  +--------+--------+---------+------+-------------+--------------+------+-------------+--------------+
  |       1|       2|   parent|     1|            1|             1|     2|         null|             2|
  |       1|       3|   parent|     1|            1|             1|     3|            2|          null|
  |       2|       4|   parent|     2|         null|             2|     4|         null|          null|
  +--------+--------+---------+------+-------------+--------------+------+-------------+--------------+
  </code></pre>
  
  <p>So far so good, now a recursive function to polyfill <code>null</code> values down the hierarchy:</p>
  
  <pre><code>def fillVertices(vertices: DataFrame, edges: DataFrame) : (DataFrame, DataFrame) = {
    val vertexAttributes = vertices.columns.filter(c =&gt; c != "id")
    val edgeAttributes = edges.columns.filter(c =&gt; (c != "src" &amp;&amp; c != "dst"))
  
    val messages = triplets(vertices,edges).select(
      Seq(col("edge_src"), col("edge_dst")) ++ vertexAttributes.map(attr =&gt; when(col("src_" + attr).isNotNull &amp;&amp; col("dst_" + attr).isNull, col("src_" + attr)) as "msg_" + attr):_*
    ).filter(
      vertexAttributes.map(attr =&gt; col("msg_" + attr).isNotNull).fold(lit(false)){ (a,b) =&gt; a || b }
    ).groupBy(col("edge_dst") as "msg_dst")
     .agg(max(col("msg_" + vertexAttributes(0))) as ("msg_" + vertexAttributes(0)), vertexAttributes.slice(1,vertexAttributes.length).map(c =&gt; max(col("msg_" + c)) as ("msg_" + c)):_*)
  
    if (! messages.rdd.isEmpty) {
      val newVerts = vertices.join(messages, col("id") === col("msg_dst"), "left_outer").select(Seq(col("id")) ++ vertexAttributes.map(c =&gt; coalesce(col(c), col("msg_" + c)) as c):_*)
      fillVertices(newVerts, edges)
    }
    else (vertices,edges)
  }
  </code></pre>
  
  <p>If you do <code>fillVertices(nodes,edges)._1.show</code> it will indeed show the correct results -- all the nodes have their <code>null</code> values correctly populated. However, it takes a ridiculous amount of computational stages.</p>
  
  <p>Again note that this is very similar behavior to what I have seen with <code>GraphFrames</code> -- I don't think it's specific to what I am doing, but rather a general problem with recursive algorithms in Spark.</p>
  
  <p>Like I said, I have tried checkpointing the underlying <code>RDD</code> and it seems to help. I use this to checkpoint a <code>DataFrame</code>:</p>
  
  <pre><code>sc.setCheckpointDir("/your/checkpoint/dir")
  def dfCheckpoint(df: DataFrame) : DataFrame = {
    df.rdd.checkpoint
    if (df.rdd.count &gt; 0) {
      df.sqlContext.createDataFrame(df.rdd, df.schema)
    }
    else df
  }
  </code></pre>
  
  <p>Then to test it side by side, here's the same algorithm as above, except that the newly created node <code>DataFrame</code> gets checkpointed before it is returned.</p>
  
  <pre><code>def fillVerticesCheckpoint(vertices: DataFrame, edges: DataFrame) : (DataFrame, DataFrame) = {
    val vertexAttributes = vertices.columns.filter(c =&gt; c != "id")
    val edgeAttributes = edges.columns.filter(c =&gt; (c != "src" &amp;&amp; c != "dst"))
  
    val messages = triplets(vertices, edges).select(
      Seq(col("edge_src"), col("edge_dst")) ++ vertexAttributes.map(attr =&gt; when(col("src_" + attr).isNotNull &amp;&amp; col("dst_" + attr).isNull, col("src_" + attr)) as "msg_" + attr):_*
    ).filter(
      vertexAttributes.map(attr =&gt; col("msg_" + attr).isNotNull).fold(lit(false)){ (a,b) =&gt; a || b }
    ).groupBy(col("edge_dst") as "msg_dst")
     .agg(max(col("msg_" + vertexAttributes(0))) as ("msg_" + vertexAttributes(0)), vertexAttributes.slice(1,vertexAttributes.length).map(c =&gt; max(col("msg_" + c)) as ("msg_" + c)):_*)
  
    if (! messages.rdd.isEmpty) {
      val newVerts = vertices.join(messages, col("id") === col("msg_dst"), "left_outer").select(Seq(col("id")) ++ vertexAttributes.map(c =&gt; coalesce(col(c), col("msg_" + c)) as c):_*)
      fillVerticesCheckpoint(dfCheckpoint(newVerts), edges)
    }
    else (vertices, edges)
  }
  </code></pre>
  
  <p>Now if you do <code>fillVerticesCheckpoint(nodes,edges)._1.show</code> it completes a lot faster. There seem to be a lot fewer stages. I don't know how to quantify it, but it seems like the checkpoint version has 1/3 the number of stages as the non checkpointed.</p>
  
  <p>Based on what I'm seeing, then, I'm guessing the answer to my first question is, yes, this is a lineage issue. The answer to my second question seems to be yes, checkpointing makes it better. But it would be great to have confirmation on these two.</p>
  
  <p>As to my last point, other ways to solve the same problem, the only think I can think of is to essentially create my own checkpoint by saving <code>DataFrames</code> to <code>Parquet</code> files between each iteration. Anybody got anything else?</p>