test: add 2PC tests for native writer

Author: Shekharrajak

spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterCommitSuite.scala

@@ -0,0 +1,335 @@
+/*
+ * 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.comet.parquet
+
+import java.io.File
+
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.execution.command.DataWritingCommandExec
+import org.apache.spark.sql.functions._
+
+import org.apache.comet.CometConf
+
+/**
+ * Test suite for Comet Native Parquet Writer commit protocol and _temporary folder handling.
+ *
+ * These tests verify that Comet properly integrates with Spark's FileCommitProtocol for reliable
+ * writes with proper 2-phase commit semantics.
+ */
+class CometParquetWriterCommitSuite extends CometTestBase {
+
+  private val nativeWriteConf = Seq(
+    CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true",
+    CometConf.COMET_EXEC_ENABLED.key -> "true",
+    CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true")
+
+  /** Helper to check if output directory contains any data files */
+  private def hasDataFiles(dir: File): Boolean = {
+    if (!dir.exists()) return false
+    dir.listFiles().exists(f => f.getName.startsWith("part-") && f.getName.endsWith(".parquet"))
+  }
+
+  // ==========================================================================
+  // Test 1: _temporary folder lifecycle - exists during write, cleaned after commit
+  // ==========================================================================
+  test("_temporary folder is created during write and cleaned up after commit") {
+    withTempPath { dir =>
+      val outputPath = new File(dir, "output").getAbsolutePath
+
+      // Create a large dataset to ensure write takes some time
+      val df = spark
+        .range(0, 100000, 1, 4)
+        .selectExpr("id", "id * 2 as value")
+
+      withSQLConf(nativeWriteConf: _*) {
+        // Start write in background thread
+        @volatile var writeStarted = false
+        @volatile var writeException: Option[Throwable] = None
+        val writeThread = new Thread(() => {
+          try {
+            writeStarted = true
+            df.write.parquet(outputPath)
+          } catch {
+            case e: Throwable => writeException = Some(e)
+          }
+        })
+        writeThread.start()
+
+        // Wait for write to start
+        CometWriteTestHelpers.waitForCondition(writeStarted, timeoutMs = 5000)
+
+        // Wait for _temporary folder to appear (indicates write started)
+        val tempExists = CometWriteTestHelpers.waitForCondition(
+          CometWriteTestHelpers.hasTemporaryFolder(outputPath),
+          timeoutMs = 10000)
+
+        // If _temporary folder appears during write, verify it exists
+        if (tempExists) {
+          assert(
+            CometWriteTestHelpers.hasTemporaryFolder(outputPath),
+            "_temporary folder should be created during write")
+
+          // Verify temp files exist
+          val tempFileCount = CometWriteTestHelpers.countTemporaryFiles(outputPath)
+          assert(tempFileCount > 0, s"Expected temp files during write, found $tempFileCount")
+        }
+
+        // Wait for write to complete
+        writeThread.join(30000)
+        assert(!writeThread.isAlive, "Write should complete within 30 seconds")
+
+        // Check if write failed
+        writeException.foreach(throw _)
+
+        // After successful commit, _temporary should be cleaned up
+        assert(
+          !CometWriteTestHelpers.hasTemporaryFolder(outputPath),
+          "_temporary folder should be cleaned up after successful commit")
+
+        // Verify final output exists
+        val outputDir = new File(outputPath)
+        assert(hasDataFiles(outputDir), "Final data files should exist")
+
+        // Verify data integrity
+        val readDf = spark.read.parquet(outputPath)
+        assert(readDf.count() == 100000, "All rows should be committed")
+      }
+    }
+  }
+
+  // ==========================================================================
+  // Test 2: _temporary cleanup on task failure
+  // ==========================================================================
+  test("_temporary folder is cleaned up on task failure") {
+    withTempPath { dir =>
+      val outputPath = new File(dir, "output").getAbsolutePath
+
+      // Create dataframe that will fail during processing
+      val divideByZero = udf((x: Long) => { x / (x - 100) })
+      val df = spark
+        .range(0, 1000, 1, 1) // Single partition to avoid race conditions
+        .select(divideByZero(col("id")).as("value"))
+
+      withSQLConf(nativeWriteConf: _*) {
+        // Write should fail with division by zero
+        intercept[Exception] {
+          df.write.parquet(outputPath)
+        }
+
+        // After failure, _temporary should be cleaned up
+        // Note: May need a small delay for cleanup to complete
+        Thread.sleep(1000)
+        assert(
+          !CometWriteTestHelpers.hasTemporaryFolder(outputPath),
+          "_temporary folder should be cleaned up after task failure")
+
+        // No data files should exist
+        val outputDir = new File(outputPath)
+        if (outputDir.exists()) {
+          assert(!hasDataFiles(outputDir), "No data files should exist after failure")
+        }
+      }
+    }
+  }
+
+  // ==========================================================================
+  // Test 3: _temporary with multiple concurrent tasks
+  // ==========================================================================
+  test("_temporary folder handles concurrent tasks correctly") {
+    withTempPath { dir =>
+      val outputPath = new File(dir, "output").getAbsolutePath
+
+      val df = spark
+        .range(0, 50000, 1, 10) // 10 partitions
+        .selectExpr("id", "id * 2 as value")
+
+      withSQLConf(nativeWriteConf: _*) {
+        // Start write in background
+        @volatile var writeStarted = false
+        @volatile var writeException: Option[Throwable] = None
+        val writeThread = new Thread(() => {
+          try {
+            writeStarted = true
+            df.write.parquet(outputPath)
+          } catch {
+            case e: Throwable => writeException = Some(e)
+          }
+        })
+        writeThread.start()
+
+        // Wait for write to start
+        CometWriteTestHelpers.waitForCondition(writeStarted, timeoutMs = 5000)
+
+        // Wait for _temporary to exist (if it does during write)
+        val tempAppeared = CometWriteTestHelpers.waitForCondition(
+          CometWriteTestHelpers.hasTemporaryFolder(outputPath),
+          timeoutMs = 10000)
+
+        // If _temporary appeared, check structure
+        if (tempAppeared) {
+          val subfolders = CometWriteTestHelpers.getTemporarySubfolders(outputPath)
+          // Should have job-level folder structure
+          assert(subfolders.nonEmpty, "Should have job tracking folders")
+        }
+
+        writeThread.join(30000)
+
+        // Check if write failed
+        writeException.foreach(throw _)
+
+        // After commit, everything cleaned up
+        assert(
+          !CometWriteTestHelpers.hasTemporaryFolder(outputPath),
+          "_temporary should be cleaned up after commit")
+
+        // Verify all partitions wrote successfully
+        val outputDir = new File(outputPath)
+        assert(hasDataFiles(outputDir), "Data files should exist")
+
+        // Verify data integrity
+        val readDf = spark.read.parquet(outputPath)
+        assert(readDf.count() == 50000, "All rows should be committed")
+      }
+    }
+  }
+
+  // ==========================================================================
+  // Test 4: _temporary with overwrite mode
+  // ==========================================================================
+  test("_temporary folder is cleaned up on overwrite") {
+    withTempPath { dir =>
+      val outputPath = new File(dir, "output").getAbsolutePath
+
+      withSQLConf(nativeWriteConf: _*) {
+        // First write
+        spark.range(1000).write.parquet(outputPath)
+        assert(
+          !CometWriteTestHelpers.hasTemporaryFolder(outputPath),
+          "_temporary should be cleaned up after first write")
+
+        // Verify first write
+        val count1 = spark.read.parquet(outputPath).count()
+        assert(count1 == 1000, "First write should have 1000 rows")
+
+        // Overwrite
+        spark.range(500).write.mode("overwrite").parquet(outputPath)
+        assert(
+          !CometWriteTestHelpers.hasTemporaryFolder(outputPath),
+          "_temporary should be cleaned up after overwrite")
+
+        // Verify final state
+        val count2 = spark.read.parquet(outputPath).count()
+        assert(count2 == 500, "Overwrite should result in 500 rows")
+      }
+    }
+  }
+
+  // ==========================================================================
+  // Test 5: No _temporary folder when write completes quickly
+  // ==========================================================================
+  test("small writes may not create visible _temporary folder") {
+    withTempPath { dir =>
+      val outputPath = new File(dir, "output").getAbsolutePath
+
+      withSQLConf(nativeWriteConf: _*) {
+        // Small dataset that completes quickly
+        spark.range(10).write.parquet(outputPath)
+
+        // After completion, _temporary should not exist
+        assert(
+          !CometWriteTestHelpers.hasTemporaryFolder(outputPath),
+          "_temporary should not exist after completion")
+
+        // Data should exist
+        val outputDir = new File(outputPath)
+        assert(hasDataFiles(outputDir), "Data files should exist")
+
+        // Verify data
+        val readDf = spark.read.parquet(outputPath)
+        assert(readDf.count() == 10, "Should have 10 rows")
+      }
+    }
+  }
+
+  // ==========================================================================
+  // Test 6: Multiple writes to different paths don't interfere
+  // ==========================================================================
+  test("multiple concurrent writes to different paths are isolated") {
+    withTempPath { dir1 =>
+      withTempPath { dir2 =>
+        val outputPath1 = new File(dir1, "output1").getAbsolutePath
+        val outputPath2 = new File(dir2, "output2").getAbsolutePath
+
+        withSQLConf(nativeWriteConf: _*) {
+          val df1 = spark.range(0, 10000, 1, 4)
+          val df2 = spark.range(10000, 20000, 1, 4)
+
+          // Start both writes
+          val thread1 = new Thread(() => df1.write.parquet(outputPath1))
+          val thread2 = new Thread(() => df2.write.parquet(outputPath2))
+
+          thread1.start()
+          thread2.start()
+
+          thread1.join(30000)
+          thread2.join(30000)
+
+          // Both should complete successfully
+          assert(!CometWriteTestHelpers.hasTemporaryFolder(outputPath1))
+          assert(!CometWriteTestHelpers.hasTemporaryFolder(outputPath2))
+
+          // Verify both outputs
+          assert(spark.read.parquet(outputPath1).count() == 10000)
+          assert(spark.read.parquet(outputPath2).count() == 10000)
+        }
+      }
+    }
+  }
+
+  // ==========================================================================
+  // Test 7: Verify no leftover _temporary folders from previous runs
+  // ==========================================================================
+  test("no stale _temporary folders from previous operations") {
+    withTempPath { dir =>
+      val outputPath = new File(dir, "output").getAbsolutePath
+
+      withSQLConf(nativeWriteConf: _*) {
+        // First write
+        spark.range(100).write.parquet(outputPath)
+        assert(!CometWriteTestHelpers.hasTemporaryFolder(outputPath))
+
+        // Second write (overwrite)
+        spark.range(200).write.mode("overwrite").parquet(outputPath)
+        assert(!CometWriteTestHelpers.hasTemporaryFolder(outputPath))
+
+        // Third write (overwrite)
+        spark.range(300).write.mode("overwrite").parquet(outputPath)
+        assert(!CometWriteTestHelpers.hasTemporaryFolder(outputPath))
+
+        // Verify final state
+        assert(spark.read.parquet(outputPath).count() == 300)
+
+        // No temporary folders should exist
+        val dirs = CometWriteTestHelpers.listDirectories(outputPath)
+        assert(!dirs.exists(_.startsWith("_temporary")), "No _temporary folders should exist")
+      }
+    }
+  }
+}