[SPARK-43281][SQL] Fix concurrent writer does not update file metrics

### What changes were proposed in this pull request?

`DynamicPartitionDataConcurrentWriter` it uses temp file path to get file status after commit task. However, the temp file has already moved to new path during commit task.

This pr calls `closeFile` before commit task.

### Why are the changes needed?

fix bug

### Does this PR introduce _any_ user-facing change?

yes, after this pr the metrics is correct

### How was this patch tested?

add test

Closes #40952 from ulysses-you/SPARK-43281.

Authored-by: ulysses-you <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
(cherry picked from commit 592e922)
Signed-off-by: Wenchen Fan <[email protected]>

Author: ulysses-you

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

@@ -159,9 +159,6 @@ class BasicWriteTaskStatsTracker(
   }
 
   override def getFinalStats(taskCommitTime: Long): WriteTaskStats = {
-    submittedFiles.foreach(updateFileStats)
-    submittedFiles.clear()
-
     // Reports bytesWritten and recordsWritten to the Spark output metrics.
     Option(TaskContext.get()).map(_.taskMetrics().outputMetrics).foreach { outputMetrics =>
       outputMetrics.setBytesWritten(numBytes)

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

@@ -427,6 +427,7 @@ class DynamicPartitionDataConcurrentWriter(
       if (status.outputWriter != null) {
         try {
           status.outputWriter.close()
+          statsTrackers.foreach(_.closeFile(status.outputWriter.path()))
         } finally {
           status.outputWriter = null
         }

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/BasicWriteTaskStatsTrackerSuite.scala

@@ -85,22 +85,23 @@ class BasicWriteTaskStatsTrackerSuite extends SparkFunSuite {
     val missing = new Path(tempDirPath, "missing")
     val tracker = new BasicWriteTaskStatsTracker(conf)
     tracker.newFile(missing.toString)
+    tracker.closeFile(missing.toString)
     assertStats(tracker, 0, 0)
   }
 
   test("Empty filename is forwarded") {
     val tracker = new BasicWriteTaskStatsTracker(conf)
     tracker.newFile("")
     intercept[IllegalArgumentException] {
-      finalStatus(tracker)
+      tracker.closeFile("")
     }
   }
 
   test("Null filename is only picked up in final status") {
     val tracker = new BasicWriteTaskStatsTracker(conf)
     tracker.newFile(null)
     intercept[IllegalArgumentException] {
-      finalStatus(tracker)
+      tracker.closeFile(null)
     }
   }
 
@@ -109,6 +110,7 @@ class BasicWriteTaskStatsTrackerSuite extends SparkFunSuite {
     val tracker = new BasicWriteTaskStatsTracker(conf)
     tracker.newFile(file.toString)
     touch(file)
+    tracker.closeFile(file.toString)
     assertStats(tracker, 1, 0)
   }
 
@@ -117,6 +119,7 @@ class BasicWriteTaskStatsTrackerSuite extends SparkFunSuite {
     val tracker = new BasicWriteTaskStatsTracker(conf)
     tracker.newFile(file.toString)
     write1(file)
+    tracker.closeFile(file.toString)
     assertStats(tracker, 1, len1)
   }
 
@@ -125,6 +128,7 @@ class BasicWriteTaskStatsTrackerSuite extends SparkFunSuite {
     val tracker = new BasicWriteTaskStatsTracker(conf)
     tracker.newFile(file.toString)
     val stream = localfs.create(file, true)
+    tracker.closeFile(file.toString)
     try {
       assertStats(tracker, 1, 0)
       stream.write(data1)
@@ -141,8 +145,10 @@ class BasicWriteTaskStatsTrackerSuite extends SparkFunSuite {
     val tracker = new BasicWriteTaskStatsTracker(conf)
     tracker.newFile(file1.toString)
     write1(file1)
+    tracker.closeFile(file1.toString)
     tracker.newFile(file2.toString)
     write2(file2)
+    tracker.closeFile(file2.toString)
     assertStats(tracker, 2, len1 + len2)
   }
 
@@ -153,10 +159,13 @@ class BasicWriteTaskStatsTrackerSuite extends SparkFunSuite {
     val tracker = new BasicWriteTaskStatsTracker(conf)
     tracker.newFile(file1.toString)
     write1(file1)
+    tracker.closeFile(file1.toString)
     tracker.newFile(file2.toString)
     write2(file2)
+    tracker.closeFile(file2.toString)
     tracker.newFile(file3.toString)
     touch(file3)
+    tracker.closeFile(file3.toString)
     assertStats(tracker, 3, len1 + len2)
   }
 
@@ -168,13 +177,16 @@ class BasicWriteTaskStatsTrackerSuite extends SparkFunSuite {
     // file 1
     tracker.newFile(file1.toString)
     write1(file1)
+    tracker.closeFile(file1.toString)
 
     // file 2 is noted, but not created
     tracker.newFile(file2.toString)
+    tracker.closeFile(file2.toString)
 
     // file 3 is noted & then created
     tracker.newFile(file3.toString)
     write2(file3)
+    tracker.closeFile(file3.toString)
 
     // the expected size is file1 + file3; only two files are reported
     // as found

sql/core/src/test/scala/org/apache/spark/sql/test/DataFrameReaderWriterSuite.scala

@@ -40,6 +40,7 @@ import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.catalyst.plans.logical.{AppendData, LogicalPlan, OverwriteByExpression}
 import org.apache.spark.sql.execution.QueryExecution
+import org.apache.spark.sql.execution.command.DataWritingCommandExec
 import org.apache.spark.sql.execution.datasources.{DataSourceUtils, HadoopFsRelation, LogicalRelation}
 import org.apache.spark.sql.execution.datasources.noop.NoopDataSource
 import org.apache.spark.sql.internal.SQLConf
@@ -1291,4 +1292,56 @@ class DataFrameReaderWriterSuite extends QueryTest with SharedSparkSession with
       }
     }
   }
+
+  test("SPARK-43281: Fix concurrent writer does not update file metrics") {
+    withTable("t") {
+      withSQLConf(SQLConf.MAX_CONCURRENT_OUTPUT_FILE_WRITERS.key -> "3",
+          SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "1") {
+        spark.sql("CREATE TABLE t(c int) USING parquet PARTITIONED BY (p String)")
+        var dataWriting: DataWritingCommandExec = null
+        val listener = new QueryExecutionListener {
+          override def onFailure(f: String, qe: QueryExecution, e: Exception): Unit = {}
+          override def onSuccess(funcName: String, qe: QueryExecution, duration: Long): Unit = {
+            qe.executedPlan match {
+              case dataWritingCommandExec: DataWritingCommandExec =>
+                dataWriting = dataWritingCommandExec
+              case _ =>
+            }
+          }
+        }
+        spark.listenerManager.register(listener)
+
+        def checkMetrics(sqlStr: String, numFiles: Int, numOutputRows: Long): Unit = {
+          sql(sqlStr)
+          sparkContext.listenerBus.waitUntilEmpty()
+          assert(dataWriting != null)
+          val metrics = dataWriting.cmd.metrics
+          assert(metrics.contains("numFiles"))
+          assert(metrics("numFiles").value == numFiles)
+          assert(metrics.contains("numOutputBytes"))
+          assert(metrics("numOutputBytes").value > 0)
+          assert(metrics.contains("numOutputRows"))
+          assert(metrics("numOutputRows").value == numOutputRows)
+        }
+
+        try {
+          // without fallback
+          checkMetrics(
+            "INSERT INTO TABLE t PARTITION(p) SELECT * FROM VALUES(1, 'a'),(2, 'a'),(1, 'b')",
+            numFiles = 2,
+            numOutputRows = 3)
+
+          // with fallback
+          checkMetrics(
+            """
+              |INSERT INTO TABLE t PARTITION(p)
+              |SELECT * FROM VALUES(1, 'a'),(2, 'b'),(1, 'c'),(2, 'd')""".stripMargin,
+            numFiles = 4,
+            numOutputRows = 4)
+        } finally {
+          spark.listenerManager.unregister(listener)
+        }
+      }
+    }
+  }
 }