wip_implement_dynamic_partition_mode

Author: coderfender

native/core/src/execution/operators/parquet_writer.rs

@@ -86,6 +86,7 @@ impl ParquetWriterExec {
         compression: CompressionCodec,
         partition_id: i32,
         column_names: Vec<String>,
+        partition_columns: Vec<String>
     ) -> Result<Self> {
         // Preserve the input's partitioning so each partition writes its own file
         let input_partitioning = input.output_partitioning().clone();

native/core/src/execution/planner.rs

@@ -1261,6 +1261,7 @@ impl PhysicalPlanner {
                     codec,
                     self.partition,
                     writer.column_names.clone(),
+                    writer.partition_columns.clone()
                 )?);
 
                 Ok((

native/proto/src/proto/operator.proto

@@ -245,6 +245,7 @@ message ParquetWriter {
   optional string job_id = 6;
   // Task attempt ID for this specific task
   optional int32 task_attempt_id = 7;
+  repeated string partition_columns = 8;  // list of partition columns to support dynamic partitioning mode
 }
 
 enum AggregateMode {

spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala

@@ -20,9 +20,7 @@
 package org.apache.comet.serde.operator
 
 import java.util.Locale
-
 import scala.jdk.CollectionConverters._
-
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 import org.apache.spark.SparkException
@@ -32,12 +30,12 @@ import org.apache.spark.sql.execution.command.DataWritingCommandExec
 import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, WriteFilesExec}
 import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
 import org.apache.spark.sql.internal.SQLConf
-
 import org.apache.comet.{CometConf, ConfigEntry, DataTypeSupport}
 import org.apache.comet.CometSparkSessionExtensions.withInfo
 import org.apache.comet.serde.{CometOperatorSerde, Incompatible, OperatorOuterClass, SupportLevel, Unsupported}
 import org.apache.comet.serde.OperatorOuterClass.Operator
 import org.apache.comet.serde.QueryPlanSerde.serializeDataType
+import org.apache.spark.sql.internal.SQLConf.PartitionOverwriteMode
 
 /**
  * CometOperatorSerde implementation for DataWritingCommandExec that converts Parquet write
@@ -64,7 +62,7 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec
             }
 
             if (cmd.partitionColumns.nonEmpty || cmd.staticPartitions.nonEmpty) {
-              return Unsupported(Some("Partitioned writes are not supported"))
+              return Incompatible(Some("Partitioned writes are not supported"))
             }
 
             if (cmd.query.output.exists(attr => DataTypeSupport.isComplexType(attr.dataType))) {
@@ -135,6 +133,7 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec
         .addAllColumnNames(cmd.query.output.map(_.name).asJava)
         // Note: work_dir, job_id, and task_attempt_id will be set at execution time
         // in CometNativeWriteExec, as they depend on the Spark task context
+        .addPartitionColumns(cmd.partitionColumns.map(_.toString()).mkString(","))
         .build()
 
       val writerOperator = Operator
@@ -159,16 +158,6 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec
     val cmd = op.cmd.asInstanceOf[InsertIntoHadoopFsRelationCommand]
     val outputPath = cmd.outputPath.toString
 
-    // SaveMode.Overwrite - delete existing output in the driver itself
-    if (cmd.mode == SaveMode.Overwrite) {
-      val outputPathObj = new Path(outputPath)
-      val fs = outputPathObj.getFileSystem(new Configuration())
-
-      if (fs.exists(outputPathObj)) {
-        fs.delete(outputPathObj, true)
-      }
-    }
-
     // Get the child plan from the WriteFilesExec or use the child directly
     val childPlan = op.child match {
       case writeFiles: WriteFilesExec =>
@@ -188,14 +177,14 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec
           classOf[org.apache.spark.sql.execution.datasources.SQLHadoopMapReduceCommitProtocol]
         val constructor =
           committerClass.getConstructor(classOf[String], classOf[String], classOf[Boolean])
-        Some(
-          constructor
-            .newInstance(
-              jobId,
-              outputPath,
-              java.lang.Boolean.FALSE // dynamicPartitionOverwrite = false for now
-            )
-            .asInstanceOf[org.apache.spark.internal.io.FileCommitProtocol])
+
+        val isDynamicOverWriteModeEnabled = cmd.partitionColumns.nonEmpty &&
+          SQLConf.get.partitionOverwriteMode == PartitionOverwriteMode.DYNAMIC
+
+          Some(
+            constructor
+              .newInstance(jobId, outputPath, isDynamicOverWriteModeEnabled)
+              .asInstanceOf[org.apache.spark.internal.io.FileCommitProtocol])
       } catch {
         case e: Exception =>
           throw new SparkException(s"Could not instantiate FileCommitProtocol: ${e.getMessage}")