Additional tests for broadcast hash join

integration_tests/src/test/scala/com/nvidia/spark/rapids/JoinsSuite.scala

@@ -17,6 +17,8 @@
 package com.nvidia.spark.rapids
 
 import org.apache.spark.SparkConf
+import org.apache.spark.sql.functions.broadcast
+import org.apache.spark.sql.rapids.execution.GpuBroadcastHashJoinExec
 
 class JoinsSuite extends SparkQueryCompareTestSuite {
 
@@ -97,4 +99,45 @@ class JoinsSuite extends SparkQueryCompareTestSuite {
     mixedDfWithNulls, mixedDfWithNulls, sortBeforeRepart = true) {
     (A, B) => A.join(B, A("longs") === B("longs"), "LeftAnti")
   }
+
+  test("broadcast hint isn't propagated after a join") {
+    val conf = new SparkConf()
+      .set("spark.sql.autoBroadcastJoinThreshold", "-1")
+
+    withGpuSparkSession(spark => {
+      val df1 = longsDf(spark)
+      val df2 = nonZeroLongsDf(spark)
+
+      val df3 = df1.join(broadcast(df2), Seq("longs"), "inner").drop(df2("longs"))
+      val df4 = longsDf(spark)
+      val df5 = df4.join(df3, Seq("longs"), "inner")
+
+      val plan = df5.queryExecution.executedPlan
+
+      assert(plan.collect { case p: GpuBroadcastHashJoinExec => p }.size === 1)
+      assert(plan.collect { case p: GpuShuffledHashJoinExec => p }.size === 1)
+    }, conf)
+  }
+
+  test("broadcast hint in SQL") {
+    withGpuSparkSession(spark => {
+      longsDf(spark).createOrReplaceTempView("t")
+      longsDf(spark).createOrReplaceTempView("u")
+
+      for (name <- Seq("BROADCAST", "BROADCASTJOIN", "MAPJOIN")) {
+        val plan1 = spark.sql(s"SELECT /*+ $name(t) */ * FROM t JOIN u ON t.longs = u.longs")
+          .queryExecution.executedPlan
+        val plan2 = spark.sql(s"SELECT /*+ $name(u) */ * FROM t JOIN u ON t.longs = u.longs")
+          .queryExecution.executedPlan
+
+        val res1 = plan1.find(_.isInstanceOf[GpuBroadcastHashJoinExec])
+        val res2 = plan2.find(_.isInstanceOf[GpuBroadcastHashJoinExec])
+
+        assert(res1.get.asInstanceOf[GpuBroadcastHashJoinExec].buildSide.toString
+          .equals("BuildLeft"))
+        assert(res2.get.asInstanceOf[GpuBroadcastHashJoinExec].buildSide.toString
+          .equals("BuildRight"))
+      }
+    })
+  }
 }