CODEWITHSUNDEEP
arraysscalaapache-spark
user2805885
user2805885

Reputation: 1705

groupBy multiple keys in scala

I have a dataframe similar to following:

val df = sc.parallelize(Seq((100, 1, 1), (100, 1,2), (100, 2,3), (200, 1,1), (200, 2,3), (200, 2, 2), (200, 3, 1), (200, 3,2), (300, 1,1), (300,1,2), (300, 2,5), (400, 1, 6))).toDF("_c0", "_c1", "_c2")

+---+---+--------------------+
|_c0|_c1|                 _c2|
+---+---+--------------------+
|100|  1|1                   |
|100|  1|2                   |
|100|  2|3                   |
|200|  1|1                   |
|200|  2|3                   |
|200|  2|2                   |
|200|  3|1                   |
|200|  3|2                   |
|300|  1|1                   |
|300|  1|2                   |
|300|  2|5                   |
|400|  1|6                   |

I need to groupBy _c0 and _c1 and get some rdd like this:

res9: Array[Array[Array[Int]]] = Array(Array(Array(1, 2), Array(3)), Array(Array(1), Array(3, 2), Array(1, 2)), Array(Array(1, 2), Array(5)), Array(Array(6)))

It is an array of array, I am new in scala. Please try to help

Upvotes: 2

Views: 1824

Answers (1)

vindev
vindev

Reputation: 2280

You can first groupBy _c0 and _c1 together and then groupBy only _c1 to get your desired result. Below is the code for same.

//first group by "_c0" and "_c1"
val res = df.groupBy("_c0", "_c1").agg(collect_list("_c2").as("_c2"))
  //group by "_c0"
  .groupBy("_c0").agg(collect_list("_c2").as("_c2"))
  .select("_c2")

res.show(false)

//output
//+---------------------------------------------------------+
//|_c2                                                      |
//+---------------------------------------------------------+
//|[WrappedArray(1, 2), WrappedArray(5)]                    |
//|[WrappedArray(1, 2), WrappedArray(3)]                    |
//|[WrappedArray(6)]                                        |
//|[WrappedArray(3, 2), WrappedArray(1, 2), WrappedArray(1)]|
//+---------------------------------------------------------+

For converting this to RDD use .rdd to the resulting dataframe. 

import scala.collection.mutable.WrappedArray
val rdd = res.rdd.map(x => x.get(0)
  .asInstanceOf[WrappedArray[WrappedArray[Int]]].array.map(x => x.toArray))

//to get the content or rdd(Don't use it if data is too big)
rdd.collect()
//output
//Array(Array(Array(1, 2), Array(5)), Array(Array(1, 2), Array(3)), Array(Array(6)), Array(Array(3, 2), Array(1, 2), Array(1)))

Upvotes: 2

Related Questions