[SPARK-56304][SQL] V2 ifPartitionNotExists support for file table INSERT INTO

Author: LuciferYang

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -1159,14 +1159,16 @@ class Analyzer(
         throw QueryCompilationErrors.unsupportedInsertReplaceOnOrUsing(
           i.table.asInstanceOf[DataSourceV2Relation].table.name())
 
-      case i: InsertIntoStatement
-          if i.table.isInstanceOf[DataSourceV2Relation] &&
-            i.query.resolved &&
-            i.replaceCriteriaOpt.isEmpty =>
-        val r = i.table.asInstanceOf[DataSourceV2Relation]
-        // ifPartitionNotExists is append with validation, but validation is not supported
+      case i @ InsertIntoStatement(r: DataSourceV2Relation, _, _, _, _, _, _, _, _)
+          if i.query.resolved && i.replaceCriteriaOpt.isEmpty =>
+        // SPARK-56304: allow ifPartitionNotExists for tables that
+        // support partition management and overwrite-by-filter
         if (i.ifPartitionNotExists) {
-          throw QueryCompilationErrors.unsupportedIfNotExistsError(r.table.name)
+          val caps = r.table.capabilities
+          if (!caps.contains(TableCapability.OVERWRITE_BY_FILTER) ||
+              !r.table.isInstanceOf[SupportsPartitionManagement]) {
+            throw QueryCompilationErrors.unsupportedIfNotExistsError(r.table.name)
+          }
         }
 
         // Create a project if this is an INSERT INTO BY NAME query.
@@ -1209,17 +1211,27 @@ class Analyzer(
               withSchemaEvolution = i.withSchemaEvolution)
           }
         } else {
+          val extraOpts = if (i.ifPartitionNotExists) {
+            Map("ifPartitionNotExists" -> "true") ++
+              staticPartitions.map { case (k, v) =>
+                s"__staticPartition.$k" -> v
+              }
+          } else {
+            Map.empty[String, String]
+          }
           if (isByName) {
             OverwriteByExpression.byName(
               table = r,
               df = query,
               deleteExpr = staticDeleteExpression(r, staticPartitions),
+              writeOptions = extraOpts,
               withSchemaEvolution = i.withSchemaEvolution)
           } else {
             OverwriteByExpression.byPosition(
               table = r,
               query = query,
               deleteExpr = staticDeleteExpression(r, staticPartitions),
+              writeOptions = extraOpts,
               withSchemaEvolution = i.withSchemaEvolution)
           }
         }

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

@@ -409,6 +409,31 @@ class DataSourceV2Strategy(session: SparkSession)
             v1, v2Write.getClass.getName, classOf[V1Write].getName)
       }
 
+    // SPARK-56304: skip write if target partition already exists
+    case OverwriteByExpression(
+        r: DataSourceV2Relation, _, query, writeOptions, _, _, Some(write), _)
+        if writeOptions.getOrElse("ifPartitionNotExists", "false") == "true"
+          && r.table.isInstanceOf[FileTable] =>
+      val ft = r.table.asInstanceOf[FileTable]
+      val prefix = "__staticPartition."
+      val staticSpec = writeOptions
+        .filter(_._1.startsWith(prefix))
+        .map { case (k, v) => k.stripPrefix(prefix) -> v }
+      // Check filesystem for partition existence
+      val partPath = ft.partitionSchema().fieldNames
+        .flatMap(col => staticSpec.get(col).map(v => s"$col=$v"))
+        .mkString("/")
+      val rootPath = ft.fileIndex.rootPaths.head
+      val hadoopConf = session.sessionState.newHadoopConf()
+      val fs = rootPath.getFileSystem(hadoopConf)
+      val targetPath = new Path(rootPath, partPath)
+      if (partPath.nonEmpty && fs.exists(targetPath)) {
+        LocalTableScanExec(Nil, Nil, None) :: Nil
+      } else {
+        OverwriteByExpressionExec(
+          planLater(query), refreshCache(r), write) :: Nil
+      }
+
     case OverwriteByExpression(
         r: DataSourceV2Relation, _, query, _, _, _, Some(write), _) =>
       OverwriteByExpressionExec(planLater(query), refreshCache(r), write) :: Nil

sql/core/src/test/scala/org/apache/spark/sql/connector/FileDataSourceV2WriteSuite.scala

@@ -723,6 +723,32 @@ class FileDataSourceV2WriteSuite extends QueryTest with SharedSparkSession {
     }
   }
 
+  test("SPARK-56304: INSERT OVERWRITE IF NOT EXISTS skips when partition exists") {
+    withTable("t") {
+      sql("CREATE TABLE t (id BIGINT, part INT)" +
+        " USING parquet PARTITIONED BY (part)")
+      sql("INSERT INTO t VALUES (1, 1), (2, 2)")
+      // IF NOT EXISTS: partition part=1 exists, should skip
+      sql("INSERT OVERWRITE TABLE t PARTITION(part=1) IF NOT EXISTS" +
+        " SELECT 999")
+      checkAnswer(
+        sql("SELECT * FROM t ORDER BY part"),
+        Seq(Row(1, 1), Row(2, 2)))
+      // Without IF NOT EXISTS: partition part=1 is replaced
+      sql("INSERT OVERWRITE TABLE t PARTITION(part=1)" +
+        " SELECT 999")
+      checkAnswer(
+        sql("SELECT * FROM t ORDER BY part"),
+        Seq(Row(999, 1), Row(2, 2)))
+      // IF NOT EXISTS: partition part=3 does not exist, should write
+      sql("INSERT OVERWRITE TABLE t PARTITION(part=3) IF NOT EXISTS" +
+        " SELECT 300")
+      checkAnswer(
+        sql("SELECT * FROM t ORDER BY part"),
+        Seq(Row(999, 1), Row(2, 2), Row(300, 3)))
+    }
+  }
+
   test("SELECT FROM format.path uses V2 path") {
     Seq("parquet", "orc", "json").foreach { format =>
       withTempPath { path =>