[VL] Honor BatchWrite.useCommitCoordinator in columnar V2 writes

gluten-substrait/src/main/scala/org/apache/gluten/execution/ColumnarV2TableWriteExec.scala

@@ -78,6 +78,7 @@ trait ColumnarV2TableWriteExec extends V2TableWriteExec with ValidatablePlan {
     // introduce a local var to avoid serializing the whole class
     val task = writingTaskBatch
     val messages = new Array[WriterCommitMessage](rdd.partitions.length)
+    val useCommitCoordinator = batchWrite.useCommitCoordinator
     val totalNumRowsAccumulator = new LongAccumulator()
 
     logInfo(
@@ -99,7 +100,7 @@ trait ColumnarV2TableWriteExec extends V2TableWriteExec with ValidatablePlan {
       sparkContext.runJob(
         rdd,
         (context: TaskContext, iter: Iterator[ColumnarBatch]) =>
-          task.run(factory, context, iter, writeMetrics),
+          task.run(factory, context, iter, useCommitCoordinator, writeMetrics),
         rdd.partitions.indices,
         (index, result: DataWritingColumnarBatchSparkTaskResult) => {
           val commitMessage = result.writerCommitMessage

gluten-substrait/src/main/scala/org/apache/spark/sql/datasources/v2/ColumnarWriteToDataSourceV2Exec.scala

@@ -18,9 +18,10 @@ package org.apache.spark.sql.datasources.v2
 
 import org.apache.gluten.connector.write.ColumnarBatchDataWriterFactory
 
-import org.apache.spark.TaskContext
+import org.apache.spark.{SparkEnv, TaskContext}
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.connector.write._
+import org.apache.spark.sql.errors.QueryExecutionErrors
 import org.apache.spark.sql.execution.datasources.v2.StreamWriterCommitProgress
 import org.apache.spark.sql.execution.metric.{CustomMetrics, SQLMetric}
 import org.apache.spark.sql.vectorized.ColumnarBatch
@@ -40,6 +41,7 @@ trait WritingColumnarBatchSparkTask[W <: DataWriter[ColumnarBatch]]
       factory: ColumnarBatchDataWriterFactory,
       context: TaskContext,
       iter: Iterator[ColumnarBatch],
+      useCommitCoordinator: Boolean,
       customMetrics: Map[String, SQLMetric]): DataWritingColumnarBatchSparkTaskResult = {
     val stageId = context.stageId()
     val stageAttempt = context.stageAttemptNumber()
@@ -48,21 +50,46 @@ trait WritingColumnarBatchSparkTask[W <: DataWriter[ColumnarBatch]]
     val attemptId = context.attemptNumber()
     val dataWriter = factory.createWriter(partId, taskId).asInstanceOf[W]
 
-    var count = 0
+    var count = 0L
+    var committed = false
     // write the data and commit this writer.
     Utils.tryWithSafeFinallyAndFailureCallbacks(block = {
       while (iter.hasNext) {
         CustomMetrics.updateMetrics(dataWriter.currentMetricsValues, customMetrics)
         val batch = iter.next()
         // Count is here.
-        count += batch.numRows()
+        count += batch.numRows().toLong
         write(dataWriter, batch)
       }
 
       CustomMetrics.updateMetrics(dataWriter.currentMetricsValues, customMetrics)
-      logInfo(s"Writer for partition ${context.partitionId()} is committing.")
-
-      val msg = dataWriter.commit()
+      val msg = if (useCommitCoordinator) {
+        val coordinator = SparkEnv.get.outputCommitCoordinator
+        val commitAuthorized = coordinator.canCommit(stageId, stageAttempt, partId, attemptId)
+        if (commitAuthorized) {
+          logInfo(
+            s"Commit authorized for partition $partId (task $taskId, attempt $attemptId, " +
+              s"stage $stageId.$stageAttempt)")
+          val commitMessage = dataWriter.commit()
+          committed = true
+          commitMessage
+        } else {
+          val commitDeniedException = QueryExecutionErrors.commitDeniedError(
+            partId,
+            taskId,
+            attemptId,
+            stageId,
+            stageAttempt)
+          logInfo(commitDeniedException.getMessage)
+          // throwing CommitDeniedException will trigger the catch block for abort
+          throw commitDeniedException
+        }
+      } else {
+        logInfo(s"Writer for partition ${context.partitionId()} is committing.")
+        val commitMessage = dataWriter.commit()
+        committed = true
+        commitMessage
+      }
       // Native write's metrics should be updated again after commit.
       CustomMetrics.updateMetrics(dataWriter.currentMetricsValues, customMetrics)
 
@@ -74,14 +101,20 @@ trait WritingColumnarBatchSparkTask[W <: DataWriter[ColumnarBatch]]
 
     })(
       catchBlock = {
-        // If there is an error, abort this writer
-        logError(
-          s"Aborting commit for partition $partId (task $taskId, attempt $attemptId, " +
-            s"stage $stageId.$stageAttempt)")
-        dataWriter.abort()
-        logError(
-          s"Aborted commit for partition $partId (task $taskId, attempt $attemptId, " +
-            s"stage $stageId.$stageAttempt)")
+        if (!committed) {
+          // If there is an error before commit has completed, abort this writer.
+          logError(
+            s"Aborting commit for partition $partId (task $taskId, attempt $attemptId, " +
+              s"stage $stageId.$stageAttempt)")
+          dataWriter.abort()
+          logError(
+            s"Aborted commit for partition $partId (task $taskId, attempt $attemptId, " +
+              s"stage $stageId.$stageAttempt)")
+        } else {
+          logError(
+            s"Skipping abort for already committed partition $partId (task $taskId, " +
+              s"attempt $attemptId, stage $stageId.$stageAttempt)")
+        }
       },
       finallyBlock = {
         dataWriter.close()

gluten-substrait/src/test/scala/org/apache/spark/sql/datasources/v2/ColumnarDataWritingSparkTaskSuite.scala

@@ -0,0 +1,255 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.datasources.v2
+
+import org.apache.gluten.connector.write.ColumnarBatchDataWriterFactory
+
+import org.apache.spark.{SparkConf, SparkEnv, SparkException, TaskContext}
+import org.apache.spark.scheduler.OutputCommitCoordinator
+import org.apache.spark.sql.connector.metric.CustomTaskMetric
+import org.apache.spark.sql.connector.write.{DataWriter, WriterCommitMessage}
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+import org.mockito.Mockito.{mock, when}
+import org.scalatest.funsuite.AnyFunSuite
+
+class ColumnarDataWritingSparkTaskSuite extends AnyFunSuite {
+
+  test("columnar write task honors Spark output commit coordinator") {
+    val stageId = 101
+    val stageAttempt = 0
+    val partitionId = 0
+    val coordinator = new RecordingOutputCommitCoordinator
+
+    withSparkEnv(coordinator) {
+      val factory = new RecordingColumnarBatchDataWriterFactory
+
+      val firstResult = DataWritingColumnarBatchSparkTask.run(
+        factory,
+        taskContext(stageId, stageAttempt, partitionId, taskAttemptId = 100L, attemptNumber = 0),
+        Iterator.empty,
+        useCommitCoordinator = true,
+        Map.empty)
+
+      assert(firstResult.writerCommitMessage == RecordingWriterCommitMessage("writer-0"))
+      assert(factory.writers.head.commits == 1)
+      assert(factory.writers.head.aborts == 0)
+      assert(coordinator.requests == Seq((stageId, stageAttempt, partitionId, 0)))
+
+      val deniedError = intercept[Exception] {
+        DataWritingColumnarBatchSparkTask.run(
+          factory,
+          taskContext(stageId, stageAttempt, partitionId, taskAttemptId = 101L, attemptNumber = 1),
+          Iterator.empty,
+          useCommitCoordinator = true,
+          Map.empty)
+      }
+
+      assert(deniedError.getMessage.contains("Commit denied"))
+      assert(factory.writers(1).commits == 0)
+      assert(factory.writers(1).aborts == 1)
+      assert(factory.writers(1).closed == 1)
+      assert(coordinator.requests == Seq(
+        (stageId, stageAttempt, partitionId, 0),
+        (stageId, stageAttempt, partitionId, 1)))
+    }
+  }
+
+  test("columnar write task counts output rows as Long") {
+    val factory = new RecordingColumnarBatchDataWriterFactory
+
+    val result = DataWritingColumnarBatchSparkTask.run(
+      factory,
+      taskContext(102, 0, partitionId = 0, taskAttemptId = 200L, attemptNumber = 0),
+      Iterator(batchWithRows(Int.MaxValue), batchWithRows(1)),
+      useCommitCoordinator = false,
+      Map.empty
+    )
+
+    assert(result.numRows == Int.MaxValue.toLong + 1L)
+    assert(factory.writers.head.commits == 1)
+    assert(factory.writers.head.aborts == 0)
+    assert(factory.writers.head.closed == 1)
+  }
+
+  test("columnar write task aborts and closes writer when commit fails") {
+    val factory = new RecordingColumnarBatchDataWriterFactory(
+      index => new RecordingColumnarBatchWriter(s"writer-$index", failOnCommit = true))
+
+    val commitError = intercept[RuntimeException] {
+      DataWritingColumnarBatchSparkTask.run(
+        factory,
+        taskContext(103, 0, partitionId = 0, taskAttemptId = 300L, attemptNumber = 0),
+        Iterator.empty,
+        useCommitCoordinator = false,
+        Map.empty)
+    }
+
+    assert(commitError.getMessage.contains("commit failed"))
+    assert(factory.writers.head.commits == 1)
+    assert(factory.writers.head.aborts == 1)
+    assert(factory.writers.head.closed == 1)
+  }
+
+  test("columnar write task does not abort committed writer after post-commit metrics failure") {
+    val factory = new RecordingColumnarBatchDataWriterFactory(
+      index => new RecordingColumnarBatchWriter(s"writer-$index", failMetricsAfterCommit = true))
+
+    val metricsError = intercept[RuntimeException] {
+      DataWritingColumnarBatchSparkTask.run(
+        factory,
+        taskContext(104, 0, partitionId = 0, taskAttemptId = 400L, attemptNumber = 0),
+        Iterator.empty,
+        useCommitCoordinator = false,
+        Map.empty)
+    }
+
+    assert(metricsError.getMessage.contains("metrics failed"))
+    assert(factory.writers.head.commits == 1)
+    assert(factory.writers.head.aborts == 0)
+    assert(factory.writers.head.closed == 1)
+  }
+
+  test("columnar write task closes denied speculative attempt when abort fails") {
+    val stageId = 105
+    val stageAttempt = 0
+    val partitionId = 0
+    val coordinator = new RecordingOutputCommitCoordinator
+
+    withSparkEnv(coordinator) {
+      val factory = new RecordingColumnarBatchDataWriterFactory(
+        index => new RecordingColumnarBatchWriter(s"writer-$index", failOnAbort = true))
+
+      val deniedError = intercept[SparkException] {
+        DataWritingColumnarBatchSparkTask.run(
+          factory,
+          taskContext(stageId, stageAttempt, partitionId, taskAttemptId = 500L, attemptNumber = 1),
+          Iterator.empty,
+          useCommitCoordinator = true,
+          Map.empty)
+      }
+
+      assert(deniedError.getMessage.contains("Commit denied"))
+      assert(factory.writers.head.commits == 0)
+      assert(factory.writers.head.aborts == 1)
+      assert(factory.writers.head.closed == 1)
+      assert(coordinator.requests == Seq((stageId, stageAttempt, partitionId, 1)))
+    }
+  }
+
+  // Mocked rather than constructed: the TaskContextImpl and SparkEnv constructor
+  // signatures differ across the Spark versions this module compiles against.
+  private def taskContext(
+      stageId: Int,
+      stageAttempt: Int,
+      partitionId: Int,
+      taskAttemptId: Long,
+      attemptNumber: Int): TaskContext = {
+    val context = mock(classOf[TaskContext])
+    when(context.stageId()).thenReturn(stageId)
+    when(context.stageAttemptNumber()).thenReturn(stageAttempt)
+    when(context.partitionId()).thenReturn(partitionId)
+    when(context.taskAttemptId()).thenReturn(taskAttemptId)
+    when(context.attemptNumber()).thenReturn(attemptNumber)
+    context
+  }
+
+  private def withSparkEnv[T](coordinator: OutputCommitCoordinator)(body: => T): T = {
+    val previous = SparkEnv.get
+    val env = mock(classOf[SparkEnv])
+    when(env.outputCommitCoordinator).thenReturn(coordinator)
+    SparkEnv.set(env)
+    try {
+      body
+    } finally {
+      SparkEnv.set(previous)
+    }
+  }
+
+  private def batchWithRows(numRows: Int): ColumnarBatch = {
+    new ColumnarBatch(Array.empty[org.apache.spark.sql.vectorized.ColumnVector], numRows)
+  }
+
+}
+
+private class RecordingOutputCommitCoordinator
+  extends OutputCommitCoordinator(new SparkConf(), true) {
+  val requests = scala.collection.mutable.ArrayBuffer.empty[(Int, Int, Int, Int)]
+
+  override def canCommit(
+      stage: Int,
+      stageAttempt: Int,
+      partition: Int,
+      attemptNumber: Int): Boolean = {
+    requests += ((stage, stageAttempt, partition, attemptNumber))
+    attemptNumber == 0
+  }
+}
+
+private class RecordingColumnarBatchDataWriterFactory(
+    writerBuilder: Int => RecordingColumnarBatchWriter =
+      index => new RecordingColumnarBatchWriter(s"writer-$index"))
+  extends ColumnarBatchDataWriterFactory {
+  val writers = scala.collection.mutable.ArrayBuffer.empty[RecordingColumnarBatchWriter]
+
+  override def createWriter(partitionId: Int, taskId: Long): DataWriter[ColumnarBatch] = {
+    val writer = writerBuilder(writers.size)
+    writers += writer
+    writer
+  }
+}
+
+private class RecordingColumnarBatchWriter(
+    id: String,
+    failOnCommit: Boolean = false,
+    failOnAbort: Boolean = false,
+    failMetricsAfterCommit: Boolean = false)
+  extends DataWriter[ColumnarBatch] {
+  var commits = 0
+  var aborts = 0
+  var closed = 0
+
+  override def write(record: ColumnarBatch): Unit = {}
+
+  override def commit(): WriterCommitMessage = {
+    commits += 1
+    if (failOnCommit) {
+      throw new RuntimeException(s"commit failed for $id")
+    }
+    RecordingWriterCommitMessage(id)
+  }
+
+  override def currentMetricsValues(): Array[CustomTaskMetric] = {
+    if (failMetricsAfterCommit && commits > 0) {
+      throw new RuntimeException(s"metrics failed for $id")
+    }
+    Array.empty
+  }
+
+  override def abort(): Unit = {
+    aborts += 1
+    if (failOnAbort) {
+      throw new RuntimeException(s"abort failed for $id")
+    }
+  }
+
+  override def close(): Unit = {
+    closed += 1
+  }
+}
+
+private case class RecordingWriterCommitMessage(id: String) extends WriterCommitMessage