Performance comparison between groupBy + join vs window func Spark

Viewed 565

These two achieve almost the same (the only diff the order of rows):

import org.apache.spark.sql.expressions._
import org.apache.spark.sql.functions._
def windowSpec = Window.partitionBy("key")
val df1 = Seq((3, "A", 5), (1, "A", 2), (3, "A", 5), (3, "B", 13)).toDF("key", "Categ1", "value")  
df1.withColumn("avg", avg("value").over(windowSpec)).show
+---+------+-----+-----------------+
|key|Categ1|value|              avg|
+---+------+-----+-----------------+
|  1|     A|    2|              2.0|
|  3|     A|    5|7.666666666666667|
|  3|     A|    5|7.666666666666667|
|  3|     B|   13|7.666666666666667|
+---+------+-----+-----------------+

and

import org.apache.spark.sql.expressions._
import org.apache.spark.sql.functions._
val df1 = Seq((3, "A", 5), (1, "A", 2), (3, "A", 5), (3, "B", 13)).toDF("key", 
val df2 = df1.groupBy("key").agg(avg("value") as "avg")
df1.join(df2,  Seq("key")).show
+---+------+-----+-----------------+
|key|Categ1|value|              avg|
+---+------+-----+-----------------+
|  3|     A|    5|7.666666666666667|
|  1|     A|    2|              2.0|
|  3|     A|    5|7.666666666666667|
|  3|     B|   13|7.666666666666667|
+---+------+-----+-----------------+

At a first glance, I thought the first one should be faster, because it avoids the join, but by experience sometimes it turns out a different conclusion.

Is there additional overhead of window func? Is there complexity difference between window func and groupBy? Thanks.

Side question
Is it impossible to ensure type safety when using window func?
Related: How use Window aggrgates on strongly typed Spark Datasets?

2 Answers

You can take a look at window's solution physical plan with explain method:

== Physical Plan ==
Window [avg(cast(value#9 as bigint)) windowspecdefinition(key#7, specifiedwindowframe(RowFrame, unboundedpreceding$(), unboundedfollowing$())) AS avg#14], [key#7]
+- *(1) Sort [key#7 ASC NULLS FIRST], false, 0
   +- Exchange hashpartitioning(key#7, 200)
      +- LocalTableScan [key#7, Categ1#8, value#9]

it will load the data, shuffle the data to get 1 key per partition then calculate the average.

for the second solution :

== Physical Plan ==
*(3) Project [key#7, Categ1#8, value#9, avg#24]
+- *(3) BroadcastHashJoin [key#7], [key#27], Inner, BuildRight
   :- LocalTableScan [key#7, Categ1#8, value#9]
   +- BroadcastExchange HashedRelationBroadcastMode(List(cast(input[0, int, false] as bigint)))
      +- *(2) HashAggregate(keys=[key#27], functions=[avg(cast(value#29 as bigint))], output=[key#27, avg#24])
         +- Exchange hashpartitioning(key#27, 200)
            +- *(1) HashAggregate(keys=[key#27], functions=[partial_avg(cast(value#29 as bigint))], output=[key#27, sum#37, count#38L])
               +- LocalTableScan [key#27, value#29]

the execution plan is way more complex, load the data, shuffle to get 1 key par partition. your data exemple is very small, so my spark use broadcastJoin instead of sortMergeJoin or ShuffleHashjoin .. bu it can cause shuffle again then join.

it will also load your df1 2 time because your data is not persisted, and the first solution seam easier to understand.

First solution is better.

  • Leaving aside that 1) no explicit setting of number of 'shuffle partitions' evident, and 2) setting AQE true or false not stated, and 3) the fact that there is a small amount of data,

    • then noting no caching (should imho not be necessary),

Then approach 1 I ran:

val dfA = spark.table("ZZZ")
import org.apache.spark.sql.expressions._
import org.apache.spark.sql.functions._
def windowSpec = Window.partitionBy("number2")
dfA.withColumn("avg", avg("number").over(windowSpec)).explain

shows:

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- Window [number#73, lit#74, number2#75, avg(_w0#85L) windowspecdefinition(number2#75, specifiedwindowframe(RowFrame, unboundedpreceding$(), unboundedfollowing$())) AS avg#80], [number2#75]
   +- Sort [number2#75 ASC NULLS FIRST], false, 0
      +- Exchange hashpartitioning(number2#75, 200), true, [id=#141]
         +- Project [number#73, lit#74, number2#75, cast(number#73 as bigint) AS _w0#85L]
            +- FileScan parquet default.zzz[number#73,lit#74,number2#75] Batched: true, DataFilters: [], Format: Parquet, Location: CatalogFileIndex[dbfs:/user/hive/warehouse/zzz], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<number:int,lit:string>

The sort is so that within partitions the same key values can be read in sequence and summed and divided by total number of key values for that key to get avg when there is a new key value and at end. This is the quickest method as it has no JOIN as in the 2nd option.

Then approach 2 I ran:

spark.conf.set("spark.sql.autoBroadcastJoinThreshold",-1)
spark.conf.set("spark.sql.adaptive.enabled",false)
val dfB = spark.table("ZZZ")
import org.apache.spark.sql.expressions._
import org.apache.spark.sql.functions._
val dfC = dfB.groupBy("number2").agg(avg("number") as "avg")
dfB.join(dfC, Seq("number2")).explain

shows:

== Physical Plan ==
*(4) Project [number2#75, number#73, lit#74, avg#500]
+- *(4) SortMergeJoin [number2#75], [number2#505], Inner
   :- Sort [number2#75 ASC NULLS FIRST], false, 0
   :  +- Exchange hashpartitioning(number2#75, 200), true, [id=#1484]
   :     +- *(1) ColumnarToRow
   :        +- FileScan parquet default.zzz[number#73,lit#74,number2#75] Batched: true, DataFilters: [], Format: Parquet, Location: InMemoryFileIndex[dbfs:/user/hive/warehouse/zzz/number2=0, dbfs:/user/hive/warehouse/zzz/number2=..., PartitionFilters: [isnotnull(number2#75)], PushedFilters: [], ReadSchema: struct<number:int,lit:string>
   +- *(3) Sort [number2#505 ASC NULLS FIRST], false, 0
      +- *(3) HashAggregate(keys=[number2#505], functions=[finalmerge_avg(merge sum#512, count#513L) AS avg(cast(number#503 as bigint))#499])
         +- Exchange hashpartitioning(number2#505, 200), true, [id=#1491]
            +- *(2) HashAggregate(keys=[number2#505], functions=[partial_avg(cast(number#503 as bigint)) AS (sum#512, count#513L)])
               +- *(2) ColumnarToRow
                  +- FileScan parquet default.zzz[number#503,number2#505] Batched: true, DataFilters: [], Format: Parquet, Location: InMemoryFileIndex[dbfs:/user/hive/warehouse/zzz/number2=0, dbfs:/user/hive/warehouse/zzz/number2=..., PartitionFilters: [isnotnull(number2#505)], PushedFilters: [], ReadSchema: struct<number:int>

This option involves a JOIN and 2 reads from data at rest. As indicated a 'self-join' and in some cases a Union with some help, should use reuse exchange according to the ticket above, but not the case. Here we have an element of 'self join' imho, but Catalyst sees it differently. This is a non-AQE approach as you can see on the settings. It is far more complex.

Conclusion:

This type of query with AQE on can adapt itself to using broadcast hash join, hence the disabling of AQE here.

First option is the way to go as it does not read from rest twice and does not need a JOIN.

Related