Problems running Spark SQL in a loop

Question

I run out of memory when I run a Spark SQL query in a loop in Scala. I use Spark 3.1.2 in local mode with 8GB driver/executor memory. I know that Spark DataFrames are not suitable for programs with iterations since they may form long lineage chains of hundreds of operations but I always assumed that this can be fixed with careful caching to break the lineage chain (which may be impossible). I understand that caching in DataFrames is lazy so it is only fully completed after an action. So, here is a simplified version of my program:

var P = spark.sql("select id, 1.0 as val from range(1000)")
val Q = spark.sql("select id, 2.0 as val from range(1000)")
val E = spark.sql("select id, 3.0 as val from range(1000)")
E.cache.count
Q.cache.count
E.createOrReplaceTempView("E")
Q.createOrReplaceTempView("Q")
for ( i <- 1 to 20 ) {
    val oldP = P
    oldP.createOrReplaceTempView("P")
    P = spark.sql("""select Q.id, P.val*Q.val*E.val as val
                     from Q join E on Q.id=E.id join P on P.id=Q.id""")
    spark.catalog.dropTempView("P")
    P.cache.count
    oldP.unpersist
}
P.count

The P.cache.count is supposed to force caching since it has an action. I have also used P.cache.write.mode("overwrite").format("noop").save() with the same results. When I run it for a large number of iterations (>10) it runs out of memory. The Spark UI shows that there are always 3 RDDs cached (E, Q, and P), so caching and unpersisting work fine. So my question is: why Spark runs out of memory for this program? In general, is there any way to break the lineage of a long chain of transformations in dataFrames, like Spark RDDs do? Thanks!

user_18032716 · Answer

My best solution so far is to dump the data into a file and read the file. This will break the lineage of the thousands of transformations throughout the loop. Let me know if you have a better solution. Caching and checkpointing don't break the lineage. So let's say that you have a loop with n steps that has cyclic dependencies, such as P=f(P), and let's say that Spark DataFrames can do up to k steps without crushing. Then, one solution is:

    val n = 1000              // number of iterations
    val k = 10                // number of loop steps that Spark can handle
    val pdir = "/tmp/cache/"  // directory to dump parquet data

    var P = spark.sql("select id, 1.0 as val from range(1000)")
    val Q = spark.sql("select id, 2.0 as val from range(1000)")
    val E = spark.sql("select id, 3.0 as val from range(1000)")
    E.createOrReplaceTempView("E")
    Q.createOrReplaceTempView("Q")
    for ( i <- 1 to n ) {
       P.createOrReplaceTempView("P")
       P = spark.sql("""select Q.id, P.val*Q.val*E.val as val
                        from Q join E on Q.id=E.id join P on P.id=Q.id""")
       spark.catalog.dropTempView("P")
       if (i%k == 0) {
         P.write.mode("overwrite").parquet(pdir+i)
         P = spark.read.parquet(pdir+i)
       }
    }
    P.count

Every k steps, it will dump P to a file and read it again. It works fine on my workstation.

Problems running Spark SQL in a loop

1 Answers