[Carmel-7244] Backport add parameters to limit the number of dynamic partitions for data source table (apache#61)

Co-authored-by: Lantao, Jin(lajin) <[email protected]>

Author: 2 people

core/src/main/scala/org/apache/spark/internal/io/FileCommitProtocol.scala

@@ -206,26 +206,46 @@ object FileCommitProtocol extends Logging {
       className: String,
       jobId: String,
       outputPath: String,
-      dynamicPartitionOverwrite: Boolean = false): FileCommitProtocol = {
+      dynamicPartitionOverwrite: Boolean = false,
+      restrictions: Map[String, _] = Map.empty): FileCommitProtocol = {
 
     logDebug(s"Creating committer $className; job $jobId; output=$outputPath;" +
       s" dynamic=$dynamicPartitionOverwrite")
     val clazz = Utils.classForName[FileCommitProtocol](className)
     // First try the constructor with arguments (jobId: String, outputPath: String,
+    // dynamicPartitionOverwrite: Boolean,
+    // restrictions: Map[String, _]).
+    // If that doesn't exist, try the one with (jobId: String, outputPath: String,
     // dynamicPartitionOverwrite: Boolean).
-    // If that doesn't exist, try the one with (jobId: string, outputPath: String).
+    // If that still doesn't exist, try the one with (jobId: string, outputPath: String).
     try {
-      val ctor = clazz.getDeclaredConstructor(classOf[String], classOf[String], classOf[Boolean])
-      logDebug("Using (String, String, Boolean) constructor")
-      ctor.newInstance(jobId, outputPath, dynamicPartitionOverwrite.asInstanceOf[java.lang.Boolean])
+      val ctor = clazz.getDeclaredConstructor(
+        classOf[String], classOf[String], classOf[Boolean],
+        classOf[Map[String, _]])
+      logDebug("Using (String, String, Boolean, Map[String, _]) constructor")
+      ctor.newInstance(
+        jobId,
+        outputPath,
+        dynamicPartitionOverwrite.asInstanceOf[java.lang.Boolean],
+        restrictions.asInstanceOf[Map[String, _]]
+      )
     } catch {
       case _: NoSuchMethodException =>
-        logDebug("Falling back to (String, String) constructor")
-        require(!dynamicPartitionOverwrite,
-          "Dynamic Partition Overwrite is enabled but" +
-            s" the committer ${className} does not have the appropriate constructor")
-        val ctor = clazz.getDeclaredConstructor(classOf[String], classOf[String])
-        ctor.newInstance(jobId, outputPath)
+        try {
+          val ctor = clazz.getDeclaredConstructor(
+            classOf[String], classOf[String], classOf[Boolean])
+          logDebug("Using (String, String, Boolean) constructor")
+          ctor.newInstance(jobId, outputPath,
+            dynamicPartitionOverwrite.asInstanceOf[java.lang.Boolean])
+        } catch {
+          case _: NoSuchMethodException =>
+            logDebug("Falling back to (String, String) constructor")
+            require(!dynamicPartitionOverwrite,
+              "Dynamic Partition Overwrite is enabled but" +
+                s" the committer ${className} does not have the appropriate constructor")
+            val ctor = clazz.getDeclaredConstructor(classOf[String], classOf[String])
+            ctor.newInstance(jobId, outputPath)
+        }
     }
   }
 

core/src/main/scala/org/apache/spark/internal/io/HadoopMapReduceCommitProtocol.scala

@@ -83,7 +83,7 @@ class HadoopMapReduceCommitProtocol(
    * [[HadoopMapReduceCommitProtocol#path]] need not be a valid [[org.apache.hadoop.fs.Path]] for
    * committers not writing to distributed file systems.
    */
-  private val hasValidPath = Try { new Path(path) }.isSuccess
+  protected val hasValidPath = Try { new Path(path) }.isSuccess
 
   /**
    * Tracks files staged by this task for absolute output paths. These outputs are not managed by
@@ -106,6 +106,8 @@ class HadoopMapReduceCommitProtocol(
    */
   @transient protected lazy val stagingDir = getStagingDir(path, jobId)
 
+  @transient private var fileCounter: Int = _
+
   protected def setupCommitter(context: TaskAttemptContext): OutputCommitter = {
     val format = context.getOutputFormatClass.getConstructor().newInstance()
     // If OutputFormat is Configurable, we should set conf to it.
@@ -136,7 +138,7 @@ class HadoopMapReduceCommitProtocol(
         new Path(Option(f.getWorkPath).map(_.toString).getOrElse(path))
       case _ => new Path(path)
     }
-
+    fileCounter += 1
     dir.map { d =>
       new Path(new Path(stagingDir, d), filename).toString
     }.getOrElse {
@@ -159,6 +161,7 @@ class HadoopMapReduceCommitProtocol(
     val tmpOutputPath = new Path(stagingDir, UUID.randomUUID().toString() + "-" + filename).toString
 
     addedAbsPathFiles(tmpOutputPath) = absOutputPath
+    fileCounter += 1
     tmpOutputPath
   }
 
@@ -192,8 +195,8 @@ class HadoopMapReduceCommitProtocol(
     committer.commitJob(jobContext)
 
     if (hasValidPath) {
-      val (allAbsPathFiles, allPartitionPaths) =
-        taskCommits.map(_.obj.asInstanceOf[(Map[String, String], Set[String])]).unzip
+      val (allAbsPathFiles, allPartitionPaths, _) =
+        taskCommits.map(_.obj.asInstanceOf[(Map[String, String], Set[String], Int)]).unzip3
       val fs = stagingDir.getFileSystem(jobContext.getConfiguration)
 
       val filesToMove = allAbsPathFiles.foldLeft(Map[String, String]())(_ ++ _)
@@ -270,14 +273,15 @@ class HadoopMapReduceCommitProtocol(
     committer.setupTask(taskContext)
     addedAbsPathFiles = mutable.Map[String, String]()
     partitionPaths = mutable.Set[String]()
+    fileCounter = 0
   }
 
   override def commitTask(taskContext: TaskAttemptContext): TaskCommitMessage = {
     val attemptId = taskContext.getTaskAttemptID
     logTrace(s"Commit task ${attemptId}")
     SparkHadoopMapRedUtil.commitTask(
       committer, taskContext, attemptId.getJobID.getId, attemptId.getTaskID.getId)
-    new TaskCommitMessage(addedAbsPathFiles.toMap -> partitionPaths.toSet)
+    new TaskCommitMessage(addedAbsPathFiles.toMap, partitionPaths.toSet, fileCounter)
   }
 
   /**

core/src/main/scala/org/apache/spark/util/Utils.scala

@@ -2987,6 +2987,11 @@ private[spark] object Utils
     resultProps
   }
 
+  @tailrec
+  def findFirstCause(t: Throwable): Throwable = {
+    if (t.getCause == null) t else findFirstCause(t.getCause)
+  }
+
   /**
    * Convert a sequence of `Path`s to a metadata string. When the length of metadata string
    * exceeds `stopAppendingThreshold`, stop appending paths for saving memory.

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -3147,6 +3147,30 @@ object SQLConf {
       .checkValues(PartitionOverwriteMode.values.map(_.toString))
       .createWithDefault(PartitionOverwriteMode.STATIC.toString)
 
+  val DYNAMIC_PARTITION_MAX_PARTITIONS =
+    buildConf("spark.carmel.sql.dynamic.partition.maxPartitions")
+      .doc("Maximum total number of dynamic partitions allowed to be created by one DML. " +
+        s"This only takes effect when ${FILE_COMMIT_PROTOCOL_CLASS.key} set to " +
+        s"org.apache.spark.sql.execution.datasources.SQLHadoopMapReduceCommitProtocol")
+      .intConf
+      .createWithDefault(Int.MaxValue)
+
+  val DYNAMIC_PARTITION_MAX_PARTITIONS_PER_TASK =
+    buildConf("spark.carmel.sql.dynamic.partition.maxPartitionsPerTask")
+      .doc("Maximum number of dynamic partitions allowed to be created per task. " +
+        s"This only takes effect when ${FILE_COMMIT_PROTOCOL_CLASS.key} set to " +
+        s"org.apache.spark.sql.execution.datasources.SQLHadoopMapReduceCommitProtocol")
+      .intConf
+      .createWithDefault(Int.MaxValue)
+
+  val DYNAMIC_PARTITION_MAX_CREATED_FILES =
+    buildConf("spark.carmel.sql.dynamic.partition.maxCreatedFiles")
+      .doc("Maximum total number of files allowed to be created in dynamic partitions write " +
+        "by one DML. This only takes effect when ${FILE_COMMIT_PROTOCOL_CLASS.key} set to " +
+        s"org.apache.spark.sql.execution.datasources.SQLHadoopMapReduceCommitProtocol")
+      .longConf
+      .createWithDefault(Long.MaxValue)
+
   object StoreAssignmentPolicy extends Enumeration {
     val ANSI, LEGACY, STRICT = Value
   }
@@ -5199,6 +5223,12 @@ class SQLConf extends Serializable with Logging with SqlApiConf {
   def partitionOverwriteMode: PartitionOverwriteMode.Value =
     PartitionOverwriteMode.withName(getConf(PARTITION_OVERWRITE_MODE))
 
+  def maxDynamicPartitions: Int = getConf(DYNAMIC_PARTITION_MAX_PARTITIONS)
+
+  def maxDynamicPartitionsPerTask: Int = getConf(DYNAMIC_PARTITION_MAX_PARTITIONS_PER_TASK)
+
+  def maxCreatedFilesInDynamicPartition: Long = getConf(DYNAMIC_PARTITION_MAX_CREATED_FILES)
+
   def storeAssignmentPolicy: StoreAssignmentPolicy.Value =
     StoreAssignmentPolicy.withName(getConf(STORE_ASSIGNMENT_POLICY))
 

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/InsertIntoHadoopFsRelationCommand.scala

@@ -30,6 +30,7 @@ import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap
 import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors}
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.command._
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.internal.SQLConf.PartitionOverwriteMode
 import org.apache.spark.sql.util.SchemaUtils
 
@@ -108,11 +109,33 @@ case class InsertIntoHadoopFsRelationCommand(
     }
 
     val jobId = java.util.UUID.randomUUID().toString
+    val parameters = CaseInsensitiveMap(options)
+
+    val partitionOverwriteMode = parameters.get("partitionOverwriteMode")
+      // scalastyle:off caselocale
+      .map(mode => PartitionOverwriteMode.withName(mode.toUpperCase))
+      // scalastyle:on caselocale
+      .getOrElse(sparkSession.sessionState.conf.partitionOverwriteMode)
+    val enableDynamicOverwrite = partitionOverwriteMode == PartitionOverwriteMode.DYNAMIC
+    // This config only makes sense when we are overwriting a partitioned dataset with dynamic
+    // partition columns.
+    val dynamicPartitionOverwrite = enableDynamicOverwrite && mode == SaveMode.Overwrite &&
+      staticPartitions.size < partitionColumns.length
+    val dynamicPartitionRestrictions: Map[String, _] = Map(
+      SQLConf.DYNAMIC_PARTITION_MAX_PARTITIONS.key ->
+        sparkSession.sessionState.conf.maxDynamicPartitions,
+      SQLConf.DYNAMIC_PARTITION_MAX_PARTITIONS_PER_TASK.key ->
+        sparkSession.sessionState.conf.maxDynamicPartitionsPerTask,
+      SQLConf.DYNAMIC_PARTITION_MAX_CREATED_FILES.key ->
+        sparkSession.sessionState.conf.maxCreatedFilesInDynamicPartition
+    )
+
     val committer = FileCommitProtocol.instantiate(
       sparkSession.sessionState.conf.fileCommitProtocolClass,
       jobId = jobId,
       outputPath = outputPath.toString,
-      dynamicPartitionOverwrite = dynamicPartitionOverwrite)
+      dynamicPartitionOverwrite = dynamicPartitionOverwrite,
+      restrictions = dynamicPartitionRestrictions)
 
     val doInsertion = if (mode == SaveMode.Append) {
       true

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SQLHadoopMapReduceCommitProtocol.scala

@@ -22,7 +22,9 @@ import org.apache.hadoop.mapreduce.{OutputCommitter, TaskAttemptContext}
 import org.apache.hadoop.mapreduce.lib.output.FileOutputCommitter
 import org.apache.parquet.hadoop.ParquetOutputCommitter
 
+import org.apache.spark.SparkException
 import org.apache.spark.internal.Logging
+import org.apache.spark.internal.io.FileCommitProtocol.TaskCommitMessage
 import org.apache.spark.internal.io.HadoopMapReduceCommitProtocol
 import org.apache.spark.sql.internal.SQLConf
 
@@ -33,10 +35,33 @@ import org.apache.spark.sql.internal.SQLConf
 class SQLHadoopMapReduceCommitProtocol(
     jobId: String,
     path: String,
-    dynamicPartitionOverwrite: Boolean = false)
+    dynamicPartitionOverwrite: Boolean = false,
+    restrictions: Map[String, Object] = Map.empty)
   extends HadoopMapReduceCommitProtocol(jobId, path, dynamicPartitionOverwrite)
     with Serializable with Logging {
 
+  private val maxDynamicPartitions = restrictions.get(
+    SQLConf.DYNAMIC_PARTITION_MAX_PARTITIONS.key) match {
+      case Some(value) => value.asInstanceOf[Int]
+      case None => Int.MaxValue
+    }
+
+  private val maxDynamicPartitionsPerTask = restrictions.get(
+    SQLConf.DYNAMIC_PARTITION_MAX_PARTITIONS_PER_TASK.key) match {
+      case Some(value) => value.asInstanceOf[Int]
+      case None => Int.MaxValue
+  }
+
+  private val maxCreatedFilesInDynamicPartition = restrictions.get(
+    SQLConf.DYNAMIC_PARTITION_MAX_CREATED_FILES.key) match {
+    case Some(value) => value.asInstanceOf[Long]
+    case None => Long.MaxValue
+  }
+
+  // They are only used in driver
+  @volatile private var totalPartitions: Set[String] = Set.empty
+  @volatile private var totalCreatedFiles: Long = 0L
+
   override protected def setupCommitter(context: TaskAttemptContext): OutputCommitter = {
     var committer = super.setupCommitter(context)
 
@@ -85,4 +110,37 @@ class SQLHadoopMapReduceCommitProtocol(
     logInfo(s"Using output committer class ${committer.getClass.getCanonicalName}")
     committer
   }
+
+  /**
+   * Called on the driver after a task commits. This can be used to access task commit messages
+   * before the job has finished. These same task commit messages will be passed to commitJob()
+   * if the entire job succeeds.
+   * Override it to check dynamic partition limitation on driver side.
+   */
+  override def onTaskCommit(taskCommit: TaskCommitMessage): Unit = {
+    logDebug(s"onTaskCommit($taskCommit)")
+    if (hasValidPath) {
+      val (_, allPartitionPaths, fileCounter) =
+        taskCommit.obj.asInstanceOf[(Map[String, String], Set[String], Int)]
+      val partitionsPerTask = allPartitionPaths.size
+      if (partitionsPerTask > maxDynamicPartitionsPerTask) {
+        throw new SparkException(s"Task tried to create $partitionsPerTask dynamic partitions," +
+          s" which is more than $maxDynamicPartitionsPerTask. To solve this" +
+          s" try to increase ${SQLConf.DYNAMIC_PARTITION_MAX_PARTITIONS_PER_TASK.key}")
+      }
+      totalPartitions ++= allPartitionPaths
+      val totalPartitionNum = totalPartitions.size
+      if (totalPartitionNum > maxDynamicPartitions) {
+        throw new SparkException(s"Total number of dynamic partitions created is" +
+          s" $totalPartitionNum, which is more than $maxDynamicPartitions." +
+          s" To solve this try to increase ${SQLConf.DYNAMIC_PARTITION_MAX_PARTITIONS.key}")
+      }
+      totalCreatedFiles += fileCounter
+      if (totalCreatedFiles > maxCreatedFilesInDynamicPartition) {
+        throw new SparkException(s"Total number of created files now is" +
+          s" $totalCreatedFiles, which exceeds $maxCreatedFilesInDynamicPartition." +
+          s" To solve this try to increase ${SQLConf.DYNAMIC_PARTITION_MAX_CREATED_FILES.key}")
+      }
+    }
+  }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileWrite.scala

@@ -59,10 +59,20 @@ trait FileWrite extends Write {
     // Hadoop Configurations are case sensitive.
     val hadoopConf = sparkSession.sessionState.newHadoopConfWithOptions(caseSensitiveMap)
     val job = getJobInstance(hadoopConf, path)
+
+    val dynamicPartitionRestrictions: Map[String, _] = Map(
+      SQLConf.DYNAMIC_PARTITION_MAX_PARTITIONS.key ->
+        sparkSession.sessionState.conf.maxDynamicPartitions,
+      SQLConf.DYNAMIC_PARTITION_MAX_PARTITIONS_PER_TASK.key ->
+        sparkSession.sessionState.conf.maxDynamicPartitionsPerTask,
+      SQLConf.DYNAMIC_PARTITION_MAX_CREATED_FILES.key ->
+        sparkSession.sessionState.conf.maxCreatedFilesInDynamicPartition
+    )
     val committer = FileCommitProtocol.instantiate(
       sparkSession.sessionState.conf.fileCommitProtocolClass,
       jobId = java.util.UUID.randomUUID().toString,
-      outputPath = paths.head)
+      outputPath = paths.head,
+      restrictions = dynamicPartitionRestrictions)
     lazy val description =
       createWriteJobDescription(sparkSession, hadoopConf, job, paths.head, options.asScala.toMap)