Using DataFrame API's anti-join operation to efficiently compare large datasets, achieving up to 2x perf improvement

[GGraziadei]

In case of not ordering comparison and primary key defined I achieved a good result comparing two data frames using DataFrame api with anti-join.
The reason of this improvement, I suppose, is related to the data volume reduction performed on each partition, which reduces a lot the shuffling cost.
I propose this enhancement for the DatasetComparer

def assertLargeDatasetContentEqualityAntiJoin(
    df1 : DataFrame,
    df2 : DataFrame
): Unit = {

  val actualCount   = df1.count()
  val expectedCount = df2.count()

  if (actualCount != expectedCount) {
    throw new Exception(s"Count mismatch: actual=$actualCount expected=$expectedCount")
  }

  val antiJoin = df1.join(df2, df1.columns.toSeq, "anti").take(5)

  if (!antiJoin.isEmpty) {
    val sampleDiff = antiJoin.mkString(", ")
    throw new Exception(s"Content mismatch! Sample differences: $sampleDiff")
  }
}

The following performance results were generated by executing the benchmark script on a MacBook M4 Pro workstation equipped with 24GB of RAM comparing actual implementation with the one which I proposed.
All reported data points represent the average across 20 independent simulations for each tested case, ensuring statistical reliability.

env1: local[8]

--- Case 1: Equal Datasets ---
Average DataFrame anti-join time: 128 ms
Average RDD join + filter time: 213 ms
Improvement ratio (DataFrame / RDD): 1.6628260180286842

--- Case 2: 50% Equal Datasets ---
Average DataFrame anti-join time: 60 ms
Average RDD join + filter time: 133 ms
Improvement ratio (DataFrame / RDD): 2.2144702697665117

--- Case 3: 25% Equal Datasets ---
Average DataFrame anti-join time: 56 ms
Average RDD join + filter time: 134 ms
Improvement ratio (DataFrame / RDD): 2.3952877296576194

--- Case 4: Completely Different Datasets ---
Average DataFrame anti-join time: 50 ms
Average RDD join + filter time: 132 ms
Improvement ratio (DataFrame / RDD): 2.6231289373658133

env2: local[4]

--- Case 1: Equal Datasets ---
Average DataFrame anti-join time: 105 ms
Average RDD join + filter time: 134 ms
Improvement ratio (DataFrame / RDD): 1.2772551467378719

--- Case 2: 50% Equal Datasets ---
Average DataFrame anti-join time: 49 ms
Average RDD join + filter time: 98 ms
Improvement ratio (DataFrame / RDD): 1.9777205624538852

--- Case 3: 25% Equal Datasets ---
Average DataFrame anti-join time: 45 ms
Average RDD join + filter time: 93 ms
Improvement ratio (DataFrame / RDD): 2.043923667592491

--- Case 4: Completely Different Datasets ---
Average DataFrame anti-join time: 41 ms
Average RDD join + filter time: 92 ms
Improvement ratio (DataFrame / RDD): 2.2148821903595595

env3: local[10]

--- Case 1: Equal Datasets ---
Average DataFrame anti-join time: 118 ms
Average RDD join + filter time: 209 ms
Improvement ratio (DataFrame / RDD): 1.7691624797983116

--- Case 2: 50% Equal Datasets ---
Average DataFrame anti-join time: 60 ms
Average RDD join + filter time: 135 ms
Improvement ratio (DataFrame / RDD): 2.2357941298503428

--- Case 3: 25% Equal Datasets ---
Average DataFrame anti-join time: 53 ms
Average RDD join + filter time: 129 ms
Improvement ratio (DataFrame / RDD): 2.3918902952305956

--- Case 4: Completely Different Datasets ---
Average DataFrame anti-join time: 45 ms
Average RDD join + filter time: 130 ms
Improvement ratio (DataFrame / RDD): 2.8310698319353857

I asked an LLM to generate a graph

this is the benchmark

import org.apache.spark.sql.{Dataset, SparkSession}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.DataFrame
import org.apache.spark.rdd.RDD
import scala.reflect.ClassTag
import java.util.concurrent.TimeUnit
import org.apache.spark.sql.Row


val spark = SparkSession.builder()
  .appName("AntiJoinVsRDDJoinBenchmark")
  .master("local[1]")
  .getOrCreate()
import spark.implicits._

val numRows = 10000
val numTests = 20

def assertLargeDatasetContentEquality[T: ClassTag](
    ds1: Dataset[T],
    ds2: Dataset[T],
    equals: (T, T) => Boolean
): Unit = {

  val ds1RDD = ds1.rdd.cache()
  val ds2RDD = ds2.rdd.cache()

  try {
    val actualCount   = ds1RDD.count()
    val expectedCount = ds2RDD.count()

    if (actualCount != expectedCount) {
      throw new Exception(s"Count mismatch: actual=$actualCount expected=$expectedCount")
    }

    val indexed1 = ds1RDD.zipWithIndex().map { case (v, i) => (i, v) }
    val indexed2 = ds2RDD.zipWithIndex().map { case (v, i) => (i, v) }

    val unequalRDD = indexed1.join(indexed2).filter { case (_, (v1, v2)) => !equals(v1, v2) }

    if (!unequalRDD.isEmpty()) {
      val sampleDiff = unequalRDD.take(5).mkString(", ")
      throw new Exception(s"Content mismatch! Sample differences: $sampleDiff")
    }

  } finally {
    ds1RDD.unpersist()
    ds2RDD.unpersist()
  }
}

def assertLargeDatasetContentEqualityAntiJoin(
    df1 : DataFrame,
    df2 : DataFrame
): Unit = {

  val actualCount   = df1.count()
  val expectedCount = df2.count()

  if (actualCount != expectedCount) {
    throw new Exception(s"Count mismatch: actual=$actualCount expected=$expectedCount")
  }

  val antiJoin = df1.join(df2, df1.columns.toSeq, "anti").take(5)

  if (!antiJoin.isEmpty) {
    val sampleDiff = antiJoin.mkString(", ")
    throw new Exception(s"Content mismatch! Sample differences: $sampleDiff")
  }
}

def simulate(caseName : String, df1 : DataFrame, df2 : DataFrame): Unit = {
  var dfTimes = Seq[Long]()
  var rddTimes = Seq[Long]()

  for (_ <- 1 to numTests) {

    // DataFrame anti-join
    val t0 = System.nanoTime()
    try{assertLargeDatasetContentEqualityAntiJoin(df1, df2)}
    catch{case e: Exception => //do nothing
    }
    val t1 = System.nanoTime()
    dfTimes = dfTimes :+ (t1 - t0)

    // RDD join + filter
    val t2 = System.nanoTime()
    try{assertLargeDatasetContentEquality[Row](df1.toDF(), df2.toDF(), (a,b) => a.equals(b))}
    catch{case e: Exception => //do nothing
    }
    val t3 = System.nanoTime()
    rddTimes = rddTimes :+ (t3 - t2)
  }

  val avgDfTime = dfTimes.sum / numTests
  val avgRddTime = rddTimes.sum / numTests

  println(s"--- $caseName ---")
  println(s"Average DataFrame anti-join time: ${TimeUnit.NANOSECONDS.toMillis(avgDfTime)} ms")
  println(s"Average RDD join + filter time: ${TimeUnit.NANOSECONDS.toMillis(avgRddTime)} ms")
  println(s"Improvement ratio (DataFrame / RDD): ${avgRddTime.toDouble / avgDfTime.toDouble}")
}

// case 1 equal datasets
val df1 = spark.range(numRows).withColumn("key", ($"id" % (numRows)).cast("int"))
val df2 = spark.range(numRows).withColumn("key", ($"id" % (numRows)).cast("int"))
simulate("Case 1: Equal Datasets", df1, df2)

// case 2 50 % equal datasets
val df3 = spark.range(numRows).withColumn("key", ($"id" % (numRows)).cast("int"))
val df4 = spark.range(numRows).withColumn("key", ($"id" % (numRows/2)).cast("int"))
simulate("Case 2: 50% Equal Datasets", df3, df4)

// case 3 25 % equal datasets
val df7 = spark.range(numRows).withColumn("key", ($"id" % (numRows)).cast("int"))
val df8 = spark.range(numRows).withColumn("key", ($"id" % (numRows/4)).cast("int"))
simulate("Case 3: 25% Equal Datasets", df7, df8)

// case 4 completely different datasets
val df5 = spark.range(numRows).withColumn("key", ($"id" % (numRows)).cast("int"))
val df6 = spark.range(numRows).withColumn("key", ($"id" % (numRows) + numRows).cast("int"))
simulate("Case 4: Completely Different Datasets", df5, df6)


spark.stop()