Add benchmark with multiple join column

mrpowers-io · Dec 10, 2024 · ba7c5cb · ba7c5cb
1 parent 37eca14
commit ba7c5cb
Showing 1 changed file with 12 additions and 7 deletions.
diff --git a/...hmarks/src/main/scala/com/github/mrpowers/spark/fast/tests/DatasetComparerBenchmark.scala b/...hmarks/src/main/scala/com/github/mrpowers/spark/fast/tests/DatasetComparerBenchmark.scala
@@ -9,7 +9,7 @@ import java.util.concurrent.TimeUnit
 import scala.util.Try
 
 private class DatasetComparerBenchmark extends DatasetComparer {
-  lazy val spark: SparkSession = {
+  def getSparkSession: SparkSession = {
     val session = SparkSession
       .builder()
       .master("local")
@@ -26,8 +26,9 @@ private class DatasetComparerBenchmark extends DatasetComparer {
   @Measurement(iterations = 10)
   @OutputTimeUnit(TimeUnit.NANOSECONDS)
   def assertLargeDatasetEqualityV2(blackHole: Blackhole): Boolean = {
-    val ds1 = spark.range(0, 1000000, 1, 8)
-    val ds3 = ds1
+    val spark = getSparkSession
+    val ds1   = spark.range(0, 1000000, 1, 8)
+    val ds3   = ds1
 
     val result = Try(assertLargeDatasetEqualityV2(ds1, ds3))
 
@@ -42,8 +43,9 @@ private class DatasetComparerBenchmark extends DatasetComparer {
   @Measurement(iterations = 10)
   @OutputTimeUnit(TimeUnit.NANOSECONDS)
   def assertLargeDatasetEqualityV2WithSinglePrimaryKey(blackHole: Blackhole): Boolean = {
-    val ds1 = spark.range(0, 1000000, 1, 8)
-    val ds3 = ds1
+    val spark = getSparkSession
+    val ds1   = spark.range(0, 1000000, 1, 8)
+    val ds3   = ds1
 
     val result = Try(assertLargeDatasetEqualityV2(ds1, ds3, primaryKeys = Seq("id")))
 
@@ -58,8 +60,9 @@ private class DatasetComparerBenchmark extends DatasetComparer {
   @Measurement(iterations = 10)
   @OutputTimeUnit(TimeUnit.NANOSECONDS)
   def assertLargeDatasetEquality(blackHole: Blackhole): Boolean = {
-    val ds1 = spark.range(0, 1000000, 1, 8)
-    val ds3 = ds1
+    val spark = getSparkSession
+    val ds1   = spark.range(0, 1000000, 1, 8)
+    val ds3   = ds1
 
     val result = Try(assertLargeDatasetEquality(ds1, ds3))
 
@@ -74,6 +77,7 @@ private class DatasetComparerBenchmark extends DatasetComparer {
   @Measurement(iterations = 10)
   @OutputTimeUnit(TimeUnit.NANOSECONDS)
   def assertLargeDatasetEqualityV2WithCompositePrimaryKey2(blackHole: Blackhole): Boolean = {
+    val spark  = getSparkSession
     val ds1    = spark.range(0, 1000000, 1, 8).withColumn("id2", col("id") + 1)
     val ds3    = ds1
     val result = Try(assertLargeDatasetEqualityV2(ds1, ds3, primaryKeys = Seq("id", "id2")))
@@ -89,6 +93,7 @@ private class DatasetComparerBenchmark extends DatasetComparer {
   @Measurement(iterations = 10)
   @OutputTimeUnit(TimeUnit.NANOSECONDS)
   def assertLargeDatasetEqualityV2WithCompositePrimaryKey3(blackHole: Blackhole): Boolean = {
+    val spark  = getSparkSession
     val ds1    = spark.range(0, 1000000, 1, 8).withColumn("id2", col("id") + 1).withColumn("id3", col("id2") + 1)
     val ds3    = ds1
     val result = Try(assertLargeDatasetEqualityV2(ds1, ds3, primaryKeys = Seq("id", "id2", "id3")))