Shuffling elements of an RDD[List[Double]] in Spark - scala

Question

In a program I am developing using Spark 2.3 in Scala, I have an RDD[List[Double]]. Every List[Double] have the same size. I can't figure out how to perform a transformation that given the RDD
[1.0, 1.5, 4.0, 3.0],
[2.3, 5.6, 3.4, 9.0],
[4.5, 2.0, 1.0, 5.7]
transform it in the RDD
[2.3, 2.0, 1.0, 3.0],
[1.0, 5.6, 4.0, 5.7],
[4.5, 1.5, 3.4, 9.0]
where every single element of the lists is swapped among them, maintaining the same position.
For example, the first element of the first list is moved to the first position of the second list, the second element of the first list is moved to the second position of the third list, and so on.
Thanks a lot.

Answer 1

One approach to shuffling column-wise would be to break down the dataset into individual single-column DataFrames each of which gets shuffled using orderBy(rand), and then piece them back together.
To join the shuffled DataFrames, RDD zipWithIndex is applied to each of them to create row-identifying ids. Note that monotonically_increasing_id won't cut it as it doesn't guarantee generating same list of ids necessary for the final join. Hence, this is rather expensive due to the required transformation between RDD and DataFrame.
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
import org.apache.spark.sql.Row
val rdd0 = sc.parallelize(Seq(
List(1.0, 1.5, 4.0, 3.0),
List(2.3, 5.6, 3.4, 9.0),
List(4.5, 2.0, 1.0, 5.7)
))
//rdd0: org.apache.spark.rdd.RDD[List[Double]] = ...
val rdd = rdd0.map{ case x: Seq[Double] => (x(0), x(1), x(2), x(3)) }
val df = rdd.toDF("c1", "c2", "c3", "c4")
val shuffledDFs = df.columns.filter(_.startsWith("c")).map{ c =>
val subDF = df.select(c)
val subRDD = subDF.orderBy(rand).rdd.zipWithIndex.map{
case (row: Row, id: Long) => Row.fromSeq(row.toSeq :+ id)
}
spark.createDataFrame( subRDD,
StructType(subDF.schema.fields :+ StructField("idx", LongType, false))
)
}
shuffledDFs.reduce( _.join(_, Seq("idx")) ).
show
// +---+---+---+---+---+
// |idx| c1| c2| c3| c4|
// +---+---+---+---+---+
// | 0|2.3|2.0|4.0|9.0|
// | 1|1.0|5.6|3.4|3.0|
// | 2|4.5|1.5|1.0|5.7|
// +---+---+---+---+---+