[SEDONA-217] Automatically broadcast small datasets.

core/src/main/java/org/apache/sedona/core/utils/SedonaConf.java

@@ -25,6 +25,7 @@
 import org.apache.sedona.core.enums.JoinSparitionDominantSide;
 import org.apache.spark.sql.RuntimeConfig;
 import org.apache.spark.sql.SparkSession;
+import org.apache.spark.util.Utils;
 import org.locationtech.jts.geom.Envelope;
 
 import java.io.Serializable;
@@ -54,6 +55,8 @@ public class SedonaConf
 
     private GridType joinGridType;
 
+    private long autoBroadcastJoinThreshold;
+
     public static SedonaConf fromActiveSession() {
         return new SedonaConf(SparkSession.active().conf());
     }
@@ -70,6 +73,11 @@ public SedonaConf(RuntimeConfig runtimeConfig)
         this.joinBuildSide = JoinBuildSide.getBuildSide(runtimeConfig.get("sedona.join.indexbuildside", "left"));
         this.joinSparitionDominantSide = JoinSparitionDominantSide.getJoinSparitionDominantSide(runtimeConfig.get("sedona.join.spatitionside", "left"));
         this.fallbackPartitionNum = Integer.parseInt(runtimeConfig.get("sedona.join.numpartition", "-1"));
+        this.autoBroadcastJoinThreshold = bytesFromString(
+                runtimeConfig.get("sedona.join.autoBroadcastJoinThreshold",
+                        runtimeConfig.get("spark.sql.autoBroadcastJoinThreshold")
+                )
+        );
     }
 
     public boolean getUseIndex()
@@ -113,6 +121,11 @@ public int getFallbackPartitionNum()
         return fallbackPartitionNum;
     }
 
+    public long getAutoBroadcastJoinThreshold()
+    {
+        return autoBroadcastJoinThreshold;
+    }
+
     public String toString()
     {
         try {
@@ -132,4 +145,12 @@ public String toString()
             return null;
         }
     }
+
+    static long bytesFromString(String str) {
+        if (str.startsWith("-")) {
+            return -1 * Utils.byteStringAsBytes(str.substring(1));
+        } else {
+            return Utils.byteStringAsBytes(str);
+        }
+    }
 }

core/src/test/java/org/apache/sedona/core/utils/SedonaConfTest.java

@@ -53,4 +53,11 @@ public void testDatasetBoundary() {
         Envelope datasetBoundary = SedonaConf.fromActiveSession().getDatasetBoundary();
         assertEquals("Env[1.0 : 2.0, 3.0 : 4.0]", datasetBoundary.toString());
     }
+
+    @Test
+    public void testBytesFromString() {
+        assertEquals(-1, SedonaConf.bytesFromString("-1"));
+        assertEquals(1024, SedonaConf.bytesFromString("1k"));
+        assertEquals(2097152, SedonaConf.bytesFromString("2MB"));
+    }
 }

docs/api/sql/Parameter.md

@@ -28,6 +28,11 @@ sparkSession.conf.set("sedona.global.index","false")
 	* Spatial index type, only valid when "sedona.global.index" is true
 	* Default: quadtree
 	* Possible values: rtree, quadtree
+* sedona.join.autoBroadcastJoinThreshold
+	* Configures the maximum size in bytes for a table that will be broadcast to all worker nodes when performing a join.
+      By setting this value to -1 automatic broadcasting can be disabled.
+	* Default: The default value is the same as spark.sql.autoBroadcastJoinThreshold
+	* Possible values: any integer with a byte suffix i.e. 10MB or 512KB
 * sedona.join.gridtype
 	* Spatial partitioning grid type for join query
 	* Default: kdbtree
@@ -43,4 +48,4 @@ sparkSession.conf.set("sedona.global.index","false")
 * sedona.join.spatitionside **(Advanced users only!)**
 	* The dominant side in spatial partitioning stage
 	* Default: left
-	* Possible values: left, right
+	* Possible values: left, right

sql/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/JoinQueryDetector.scala

@@ -78,12 +78,28 @@ class JoinQueryDetector(sparkSession: SparkSession) extends Strategy {
     }
 
   def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
-    case Join(left, right, joinType, condition, JoinHint(leftHint, rightHint)) => { // SPARK3 anchor
-//    case Join(left, right, Inner, condition) => { // SPARK2 anchor
-      val broadcastLeft = leftHint.exists(_.strategy.contains(BROADCAST)) // SPARK3 anchor
-      val broadcastRight = rightHint.exists(_.strategy.contains(BROADCAST)) // SPARK3 anchor
-//      val broadcastLeft = left.isInstanceOf[ResolvedHint] && left.asInstanceOf[ResolvedHint].hints.broadcast  // SPARK2 anchor
-//      val broadcastRight = right.isInstanceOf[ResolvedHint] && right.asInstanceOf[ResolvedHint].hints.broadcast // SPARK2 anchor
+    case Join(left, right, joinType, condition, JoinHint(leftHint, rightHint)) => {
+      var broadcastLeft = leftHint.exists(_.strategy.contains(BROADCAST))
+      var broadcastRight = rightHint.exists(_.strategy.contains(BROADCAST))
+
+      /*
+      If either side is small we can automatically broadcast just like Spark does.
+      This only applies to inner joins as there are no optimized fallback plan for other join types.
+      It's better that users are explicit about broadcasting for other join types than seeing wildly different behavior
+      depending on data size.
+       */
+      if (!broadcastLeft && !broadcastRight && joinType == Inner) {
+        val canAutoBroadCastLeft = canAutoBroadcastBySize(left)
+        val canAutoBroadCastRight = canAutoBroadcastBySize(right)
+        if (canAutoBroadCastLeft && canAutoBroadCastRight) {
+          // Both sides can be broadcasted. Choose the smallest side.
+          broadcastLeft = left.stats.sizeInBytes <= right.stats.sizeInBytes
+          broadcastRight = !broadcastLeft
+        } else {
+          broadcastLeft = canAutoBroadCastLeft
+          broadcastRight = canAutoBroadCastRight
+        }
+      }
 
       val queryDetection: Option[JoinQueryDetection] = condition match {
         case Some(predicate: ST_Predicate) =>
@@ -137,6 +153,9 @@ class JoinQueryDetector(sparkSession: SparkSession) extends Strategy {
       Nil
   }
 
+  private def canAutoBroadcastBySize(plan: LogicalPlan) =
+    plan.stats.sizeInBytes != 0 && plan.stats.sizeInBytes <= SedonaConf.fromActiveSession.getAutoBroadcastJoinThreshold
+
   /**
     * Returns true if specified expression has at least one reference and all its references
     * map to the output of the specified plan.

sql/src/test/scala/org/apache/sedona/sql/BroadcastIndexJoinSuite.scala

@@ -1408,4 +1408,26 @@ class BroadcastIndexJoinSuite extends TestBaseScala {
     }
   }
 
+  describe("Sedona SQL automatic broadcast") {
+    it("Datasets smaller than threshold should be broadcasted") {
+      val polygonDf = buildPolygonDf.repartition(3).alias("polygon")
+      val pointDf = buildPointDf.repartition(5).alias("point")
+      val df = polygonDf.join(pointDf, expr("ST_Contains(polygon.polygonshape, point.pointshape)"))
+      sparkSession.conf.set("sedona.global.index", "true")
+      sparkSession.conf.set("sedona.join.autoBroadcastJoinThreshold", "10mb")
+
+      assert(df.queryExecution.sparkPlan.collect { case p: BroadcastIndexJoinExec => p }.size === 1)
+      sparkSession.conf.set("sedona.join.autoBroadcastJoinThreshold", "-1")
+    }
+
+    it("Datasets larger than threshold should not be broadcasted") {
+      val polygonDf = buildPolygonDf.repartition(3).alias("polygon")
+      val pointDf = buildPointDf.repartition(5).alias("point")
+      val df = polygonDf.join(pointDf, expr("ST_Contains(polygon.polygonshape, point.pointshape)"))
+      sparkSession.conf.set("sedona.global.index", "true")
+      sparkSession.conf.set("sedona.join.autoBroadcastJoinThreshold", "-1")
+
+      assert(df.queryExecution.sparkPlan.collect { case p: BroadcastIndexJoinExec => p }.size === 0)
+    }
+  }
 }

sql/src/test/scala/org/apache/sedona/sql/TestBaseScala.scala

@@ -37,6 +37,8 @@ trait TestBaseScala extends FunSpec with BeforeAndAfterAll {
     config("spark.kryo.registrator", classOf[SedonaKryoRegistrator].getName).
     master("local[*]").appName("sedonasqlScalaTest")
     .config("spark.sql.warehouse.dir", warehouseLocation)
+    // We need to be explicit about broadcasting in tests.
+    .config("sedona.join.autoBroadcastJoinThreshold", "-1")
     .getOrCreate()
 
   val resourceFolder = System.getProperty("user.dir") + "/../core/src/test/resources/"