Spark DataFrame: Is order of withColumn guaranteed? - scala

Question

Given the following code:
dataFrame
.withColumn("A", myUdf1($"x")) // withColumn1 from x
.withColumn("B", myUdf2($"y")) // withColumn2 from y
Is it guaranteed that withColumn1 will execute before withColumn2?
A better example:
dataFrame
.withColumn("A", myUdf1($"x")) // withColumn1 from x
.withColumn("B", myUdf2($"A")) // withColumn2 from A!!
Note that withColumn2 operates on A that is calculated from withColumn1.
I'm asking because I'm having inconsistent results over multiple runs of the same code and I started to think that this could be the source of the issue.
EDIT: Added more detailed code sample
val result = dataFrame
.groupBy("key")
.agg(
collect_list($"itemList").as("A"), // all items
collect_list(when($"click".isNotNull, $"itemList")).as("B") // subset of A
)
// create sparse item vector from all list of items A
.withColumn("vectorA", aggToSparseUdf($"A"))
// create sparse item vector from all list of items B (subset of A)
.withColumn("vectorB", aggToSparseUdf($"B"))
// calculate ratio vector B / A
.withColumn("ratio", divideVectors($"vectorB", $"vectorA"))
val keys: Seq[String] = result.head.getAs[Seq[String]]("key")
val values: Seq[SparseVector] = result.head.getAs[Seq[SparseVector]]("ratio")

Answer 1

It IS guaranteed that for each specific record in dataFrame, myUdf1 will be applied before myUdf2; However:
It is NOT guaranteed that myUdf1 will be applied to all records of dataFrame before myUdf2 is applied to any record - in other words, myUdf2 might be applied to some records before myUdf1 has been applied to other records
This is true because Spark would likely combine both operations together into a single stage, and execute this stage (applying myUdf1 and myUdf2) on each record of each partition.
This shouldn't pose any problem if your UDFs are "purely functional", or "idempotent", or cause no side effects - and they should be, because Spark assumes all transformations are such. If they weren't, Spark wouldn't be able to optimize execution by "combining" transformations together, running them in parallel on different partitions, retrying transformations etc.
EDIT: if you want to force UDF1 to be completely applied before applying UDF2 to any record, you'd have to force them into separate stages - this can be done, for example, by repartitioning the DataFrame:
// sample data:
val dataFrame = Seq("A", "V", "D").toDF("x")
// two UDFs with "side effects" (printing to console):
val myUdf1 = udf[String, String](x => {
println("In UDF1")
x.toLowerCase
})
val myUdf2 = udf[String, String](x => {
println("In UDF2")
x.toUpperCase
})
// repartitioning between UDFs
dataFrame
.withColumn("A", myUdf1($"x"))
.repartition(dataFrame.rdd.partitions.length + 1)
.withColumn("B", myUdf2($"A"))
.show()
// prints:
// In UDF1
// In UDF1
// In UDF1
// In UDF2
// In UDF2
// In UDF2
NOTE that this isn't bullet-proof either - if, for example, there are failures and retries, order can be once again non deterministic.