[CARMEL-7185] Backport add max split limitations for cartesian rdd (apache#29)

Author: wakun

core/src/main/scala/org/apache/spark/internal/config/package.scala

@@ -2564,4 +2564,11 @@ package object config {
       .version("3.5.0")
       .stringConf
       .createOptional
+
+  private[spark] val RDD_MAX_PARTITIONS = ConfigBuilder("spark.carmel.stage.max.partitions")
+    .doc("The maximum number of partitions allowed when constructing a RDD.")
+    .version("3.5.0")
+    .intConf
+    .checkValue(_ > 0, "The maximum number of partitions must be greater than 0.")
+    .createWithDefault(Int.MaxValue)
 }

core/src/main/scala/org/apache/spark/rdd/CartesianRDD.scala

@@ -22,6 +22,7 @@ import java.io.{IOException, ObjectOutputStream}
 import scala.reflect.ClassTag
 
 import org.apache.spark._
+import org.apache.spark.internal.config.RDD_MAX_PARTITIONS
 import org.apache.spark.util.Utils
 
 private[spark]
@@ -55,8 +56,14 @@ class CartesianRDD[T: ClassTag, U: ClassTag](
 
   val numPartitionsInRdd2 = rdd2.partitions.length
 
+  val MAX_PARTITIONS_IN_STAGE = conf.get(RDD_MAX_PARTITIONS)
+
   override def getPartitions: Array[Partition] = {
     // create the cross product split
+    if (rdd1.partitions.length * rdd2.partitions.length > MAX_PARTITIONS_IN_STAGE) {
+      throw new SparkException(s"Cartesian operation's partitions have " +
+        s"reached max limitation $MAX_PARTITIONS_IN_STAGE")
+    }
     val array = new Array[Partition](rdd1.partitions.length * rdd2.partitions.length)
     for (s1 <- rdd1.partitions; s2 <- rdd2.partitions) {
       val idx = s1.index * numPartitionsInRdd2 + s2.index