Improving performance of distinct + groupByKey on Spark - scala

Question

I am trying to learn spark and came up with this problem but my solution doesn't seem to be performing well. I was hoping someone can educate me on how I can improve the performance. The problem I have is as follows.
I have a few million tuples (e.g. (A, B), (A, C), (B, C), etc.) with possibly duplicate tuples (keys and value). What I would like to do is group up the tuples by key AND, to make it more interesting, limit the length of the grouped values to some arbitrary number (let's say 3).
So, for example, if I have:
[(A, B), (A, C), (A, D), (A, E), (B, C)]
I would expect the output to be:
[(A, [B, C, D]), (A, [E]), (B, [C]))
If any of the values for the list got longer than 3, then it would split it and you the same key listed multiple times as shown above with (A, [E]). Hopefully this makes sense.
The solution I came up with was:
val myTuples: Array[(String, String)] = ...
sparkContext.parallelize(myTuples)
.distinct() // to delete duplicates
.groupByKey() // to group up the tuples by key
.flatMapValues(values => values.grouped(3)) // split up values in groups of 3
.repartition(sparkContext.defaultParallelism)
.collect()
My solution works okay but is there a more efficient way to do this? I hear that groupByKey is very inefficient. Any help would be greatly appreciated.
Also, is there a good number I should choose for partitions? I noticed that distinct takes in a partition parameter but not sure what I should have put.
Thanks!

Answer 1

You need to re-formulate your problem slightly, as you aren't actually grouping by a single key here; in your example above, you output multiple rows for "A". In the below, I add a column which we can use additionally to group by (it will increment every 3 records), and collect_list which is a Spark SQL function to produce the arrays you are looking for. Note that by sticking to entirely Spark SQL, you get many optimisations from Spark (through "catalyst" which is a query optimiser).
import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions.Window
val data = List(("A", "B"), ("A", "C"), ("A", "D"), ("A", "E"), ("B", "C")).toDF("key","value")
val data2 = data.withColumn("index", floor(
(row_number().over(Window.partitionBy("key").orderBy("value"))-1)/3)
)
data2.show
+---+-----+-----+
|key|value|index|
+---+-----+-----+
| B| C| 0|
| A| B| 0|
| A| C| 0|
| A| D| 0|
| A| E| 1|
+---+-----+-----+
data2.groupBy("key","index").agg(collect_list("value")).show
+---+-----+-------------------+
|key|index|collect_list(value)|
+---+-----+-------------------+
| B| 0| [C]|
| A| 0| [B, C, D]|
| A| 1| [E]|
+---+-----+-------------------+