What are the Scala Spark performance of groupBy vs pivot?

Question

I am facing an issue that I have to pivot a Spark Dataframe with different aggregation functions, based on the column value I decide to pivot. I am using this other question on SO as my starting point.

Let's take the following as starting point:

scala> val data = Seq((1, "k1", "measureA", 2), (1, "k1", "measureA", 4), (1, "k1", "measureB", 5), (1, "k1", "measureB", 7), (1, "k1", "measureC", 7), (1, "k1", "measureC", 1), (2, "k1", "measureB", 8), (2, "k1", "measureC", 9), (2, "k2", "measureA", 5), (2, "k2", "measureC", 5), (2, "k2", "measureC", 8))
data: Seq[(Int, String, String, Int)] = List((1,k1,measureA,2), (1,k1,measureA,4), (1,k1,measureB,5), (1,k1,measureB,7), (1,k1,measureC,7), (1,k1,measureC,1), (2,k1,measureB,8), (2,k1,measureC,9), (2,k2,measureA,5), (2,k2,measureC,5), (2,k2,measureC,8))


scala> val df = data.toDF("ts","key","measure_type","value")
df: org.apache.spark.sql.DataFrame = [ts: int, key: string ... 2 more fields]

scala> df.show
+---+---+------------+-----+
| ts|key|measure_type|value|
+---+---+------------+-----+
|  1| k1|    measureA|    2|
|  1| k1|    measureA|    4|
|  1| k1|    measureB|    5|
|  1| k1|    measureB|    7|
|  1| k1|    measureC|    7|
|  1| k1|    measureC|    1|
|  2| k1|    measureB|    8|
|  2| k1|    measureC|    9|
|  2| k2|    measureA|    5|
|  2| k2|    measureC|    5|
|  2| k2|    measureC|    8|
+---+---+------------+-----+

What does perform better? A groupBy + agg:

val ddf = df.groupBy("ts", "key").agg(
 sum(when(col("measure_type") === "measureA",col("value"))).as("measureA"),
 avg(when(col("measure_type") === "measureB",col("value"))).as("measureB"),
 max(when(col("measure_type") === "measureC",col("value"))).as("measureC"))

ddf.show

+---+---+--------+--------+--------+
| ts|key|measureA|measureB|measureC|
+---+---+--------+--------+--------+
|  1| k1|       6|     6.0|       7|
|  2| k1|    null|     8.0|       9|
|  2| k2|       5|    null|       8|
+---+---+--------+--------+--------+

Or pivot + agg:

val listA = Seq("measureA")
val listB = Seq("measureB")
val listC = Seq("measureC")

val ddf = df.groupBy("ts", "key").pivot(col("measure_type"), Seq("measureA", "measureB", "measureC")).agg(
 sum(when(col("measure_type").isInCollection(listA),col("value"))).as("measureA"),
 avg(when(col("measure_type").isInCollection(listB),col("value"))).as("measureB"),
 max(when(col("measure_type").isInCollection(listC),col("value"))).as("measureC"))

ddf.show()

+---+---+-----------------+-----------------+-----------------+-----------------+-----------------+-----------------+-----------------+-----------------+-----------------+
| ts|key|measureA_measureA|measureA_measureB|measureA_measureC|measureB_measureA|measureB_measureB|measureB_measureC|measureC_measureA|measureC_measureB|measureC_measureC|
+---+---+-----------------+-----------------+-----------------+-----------------+-----------------+-----------------+-----------------+-----------------+-----------------+
|  2| k2|                5|             null|             null|             null|             null|             null|             null|             null|                8|
|  2| k1|             null|             null|             null|             null|              8.0|             null|             null|             null|                9|
|  1| k1|                6|             null|             null|             null|              6.0|             null|             null|             null|                7|
+---+---+-----------------+-----------------+-----------------+-----------------+-----------------+-----------------+-----------------+-----------------+-----------------+

I am aware that the second DataFrame that I got as output is different because it contains all the combinations from the Seq() of column names I passed to the pivot method and the number of the different aggregation functions (SQL CASEs, since there is when()) I decided to choose. So 3 * 3 = 9. But if you filter this second DataFrame removing the columns with only null then the result is the same.

Also I am wondering if I am doing something wrong in the second approach or there is a way to better rename the columns in order to avoid from the start the null columns.

  val ddf = df.groupBy("ts", "key").pivot(col("measure_type"), Seq("measureA", "measureB", "measureC")).agg(
     sum(when(col("measure_type").isInCollection(listA),col("value"))),
     avg(when(col("measure_type").isInCollection(listB),col("value"))),
     max(when(col("measure_type").isInCollection(listC),col("value"))))
    
    ddf:org.apache.spark.sql.DataFrame
    ts:integer
    key:string
    measureA_sum(CASE WHEN (measure_type IN (measureA)) THEN value END):long
    measureA_avg(CASE WHEN (measure_type IN (measureB)) THEN value END):double
    measureA_max(CASE WHEN (measure_type IN (measureC)) THEN value END):integer
    measureB_sum(CASE WHEN (measure_type IN (measureA)) THEN value END):long
    measureB_avg(CASE WHEN (measure_type IN (measureB)) THEN value END):double
    measureB_max(CASE WHEN (measure_type IN (measureC)) THEN value END):integer
    measureC_sum(CASE WHEN (measure_type IN (measureA)) THEN value END):long
    measureC_avg(CASE WHEN (measure_type IN (measureB)) THEN value END):double
    measureC_max(CASE WHEN (measure_type IN (measureC)) THEN value END):integer

I have decided not to post the ddf.show because of the very verbose headers. The result is the same as the pivot + agg example, just with the headers above listed.

What are the Scala Spark performance of groupBy vs pivot?

0 Answers