diff --git a/docs/source/user-guide/latest/installation.md b/docs/source/user-guide/latest/installation.md index 9fc350c25f..1e59034f7a 100644 --- a/docs/source/user-guide/latest/installation.md +++ b/docs/source/user-guide/latest/installation.md @@ -142,8 +142,8 @@ Comet will log output similar to: ```shell INFO core/src/lib.rs: Comet native library version $COMET_VERSION initialized WARN CometExecRule: Comet cannot execute some parts of this plan natively (set spark.comet.explain.fallback.enabled=false to disable this logging): - Execute InsertIntoHadoopFsRelationCommand [COMET: Native support for operator DataWritingCommandExec is disabled. Set spark.comet.parquet.write.enabled=true to enable it.] -+- WriteFiles + Execute InsertIntoHadoopFsRelationCommand ++- WriteFiles [COMET: Native support for operator WriteFilesExec is disabled. Set spark.comet.parquet.write.enabled=true to enable it.] +- LocalTableScan [COMET: Native support for operator LocalTableScanExec is disabled. Set spark.comet.exec.localTableScan.enabled=true to enable it.] ``` diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 42235b1367..4ccee468c0 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -116,9 +116,9 @@ omitted from the tables below and may be reconsidered based on demand: ## Writes -| Operator | Status | Notes | -| ------------------------ | ------ | ----------------------------------------------------------------- | -| `DataWritingCommandExec` | ⚠️ | Experimental native Parquet writes, disabled by default (opt-in). | +| Operator | Status | Notes | +| ---------------- | ------ | --------------------------------------------------------------------------------------------------------------------------------- | +| `WriteFilesExec` | ⚠️ | Experimental native Parquet writes, disabled by default (opt-in). Requires Spark 4.0+. Non-partitioned, non-bucketed writes only. | ## Python and UDF diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index dbbee713ae..223e4463f2 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -218,18 +218,11 @@ impl ParquetWriter { pub struct ParquetWriterExec { /// Input execution plan input: Arc, - /// Output file path (final destination) + /// Path of the file to write. Chosen by the JVM commit protocol (see CometWriteFilesExec) and + /// used verbatim - this operator never derives file names of its own. output_path: String, - /// Working directory for temporary files (used by FileCommitProtocol) - work_dir: String, - /// Job ID for tracking this write operation - job_id: Option, - /// Task attempt ID for this specific task - task_attempt_id: Option, /// Compression codec compression: ParquetCompression, - /// Partition ID (from Spark TaskContext) - partition_id: i32, /// Column names to use in the output Parquet file column_names: Vec, /// Object store configuration options @@ -242,15 +235,10 @@ pub struct ParquetWriterExec { impl ParquetWriterExec { /// Create a new ParquetWriterExec - #[allow(clippy::too_many_arguments)] pub fn try_new( input: Arc, output_path: String, - work_dir: String, - job_id: Option, - task_attempt_id: Option, compression: ParquetCompression, - partition_id: i32, column_names: Vec, object_store_options: HashMap, ) -> Result { @@ -267,11 +255,7 @@ impl ParquetWriterExec { Ok(ParquetWriterExec { input, output_path, - work_dir, - job_id, - task_attempt_id, compression, - partition_id, column_names, object_store_options, metrics: ExecutionPlanMetricsSet::new(), @@ -447,11 +431,7 @@ impl ExecutionPlan for ParquetWriterExec { 1 => Ok(Arc::new(ParquetWriterExec::try_new( Arc::clone(&children[0]), self.output_path.clone(), - self.work_dir.clone(), - self.job_id.clone(), - self.task_attempt_id, self.compression.clone(), - self.partition_id, self.column_names.clone(), self.object_store_options.clone(), )?)), @@ -476,8 +456,6 @@ impl ExecutionPlan for ParquetWriterExec { let runtime_env = context.runtime_env(); let input = self.input.execute(partition, context)?; let input_schema = self.input.schema(); - let work_dir = self.work_dir.clone(); - let task_attempt_id = self.task_attempt_id; let compression = self.compression.to_parquet()?; let column_names = self.column_names.clone(); @@ -492,16 +470,8 @@ impl ExecutionPlan for ParquetWriterExec { .collect(); let output_schema = Arc::new(arrow::datatypes::Schema::new(fields)); - // Generate part file name for this partition - // If using FileCommitProtocol (work_dir is set), include task_attempt_id in the filename - let part_file = if let Some(attempt_id) = task_attempt_id { - format!( - "{}/part-{:05}-{:05}.parquet", - work_dir, self.partition_id, attempt_id - ) - } else { - format!("{}/part-{:05}.parquet", work_dir, self.partition_id) - }; + // The JVM commit protocol already decided where this task writes. + let part_file = self.output_path.clone(); // Configure writer properties let props = WriterProperties::builder() @@ -851,18 +821,14 @@ mod tests { let memory_exec = Arc::new(DataSourceExec::new(Arc::new(memory_source_config))); // Create ParquetWriterExec with DataSourceExec as input - let output_path = "unused".to_string(); - let work_dir = "hdfs://namenode:9000/user/test_parquet_writer_exec".to_string(); + let output_path = + "hdfs://namenode:9000/user/test_parquet_writer_exec/part-00000.parquet".to_string(); let column_names = vec!["id".to_string(), "name".to_string()]; let parquet_writer = ParquetWriterExec::try_new( memory_exec, output_path, - work_dir, - None, // job_id - Some(123), // task_attempt_id ParquetCompression::None, - 0, // partition_id column_names, HashMap::new(), // object_store_options )?; diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 85851b17ca..cf13a67928 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1837,15 +1837,7 @@ impl PhysicalPlanner { let parquet_writer = Arc::new(ParquetWriterExec::try_new( Arc::clone(&child.native_plan), writer.output_path.clone(), - writer - .work_dir - .as_ref() - .expect("work_dir is provided") - .clone(), - writer.job_id.clone(), - writer.task_attempt_id, codec, - self.partition, writer.column_names.clone(), object_store_options, )?); diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index ced87262f3..61fd4c3f8e 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -458,16 +458,17 @@ message ShuffleWriter { } message ParquetWriter { + // Fully-qualified path of the Parquet file that this task must write, set per task by + // CometWriteFilesExec from FileCommitProtocol.newTaskTempFile. Naming and staging are owned by + // Spark's commit protocol so that task-attempt isolation, speculative execution, and committers + // that track individual files (S3A magic, streaming manifest) all behave as they do for Spark's + // own writer. The native writer uses this path verbatim. string output_path = 1; CompressionCodec compression = 2; repeated string column_names = 4; - // Working directory for temporary files (used by FileCommitProtocol) - // If not set, files are written directly to output_path - optional string work_dir = 5; - // Job ID for tracking this write operation - optional string job_id = 6; - // Task attempt ID for this specific task - optional int32 task_attempt_id = 7; + // Was work_dir / job_id / task_attempt_id, used when the native writer derived its own file + // names from the task context. File naming now comes from the commit protocol instead. + reserved 5, 6, 7; // Options for configuring object stores such as AWS S3, GCS, etc. The key-value pairs are taken // from Hadoop configuration for compatibility with Hadoop FileSystem implementations of object // stores. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 1bb524349d..3dae1c02e1 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -870,8 +870,13 @@ object CometConf extends ShimCometConf { .booleanConf .createWithEnvVarOrDefault("ENABLE_COMET_STRICT_TESTING", false) - val COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT: ConfigEntry[Boolean] = - createOperatorIncompatConfig("DataWritingCommandExec") + val COMET_OPERATOR_WRITE_FILES_ALLOW_INCOMPAT: ConfigEntry[Boolean] = + // Native writes used to replace the whole DataWritingCommandExec; they now replace only + // WriteFilesExec, so the opt-in moved with it. Keep the old key working for anyone who had + // already enabled the experimental writer. + createOperatorIncompatConfig( + "WriteFilesExec", + Some(getOperatorAllowIncompatConfigKey("DataWritingCommandExec"))) /** Create a config to enable a specific operator */ private def createExecEnabledConfig( @@ -896,15 +901,18 @@ object CometConf extends ShimCometConf { private def configKeyToEnvVar(configKey: String): String = configKey.toUpperCase(Locale.ROOT).replace('.', '_') - private def createOperatorIncompatConfig(name: String): ConfigEntry[Boolean] = { + private def createOperatorIncompatConfig( + name: String, + alternative: Option[String] = None): ConfigEntry[Boolean] = { val configKey = getOperatorAllowIncompatConfigKey(name) val envVar = configKeyToEnvVar(configKey) - conf(configKey) + val builder = conf(configKey) .category(CATEGORY_EXEC) .doc(s"Whether to allow incompatibility for operator: $name. " + s"False by default. Can be overridden with $envVar env variable") - .booleanConf - .createWithEnvVarOrDefault(envVar, false) + // ConfigBuilder mutates in place, so the result of withAlternative is `builder` itself. + alternative.foreach(builder.withAlternative(_)) + builder.booleanConf.createWithEnvVarOrDefault(envVar, false) } def isExprEnabled(name: String, conf: SQLConf = SQLConf.get): Boolean = { diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index ef2f37371c..fa6d4cfb65 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -35,7 +35,7 @@ import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec} -import org.apache.spark.sql.execution.datasources.WriteFilesExec +import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, WriteFilesExec} import org.apache.spark.sql.execution.datasources.csv.CSVFileFormat import org.apache.spark.sql.execution.datasources.json.JsonFileFormat import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat @@ -101,6 +101,14 @@ object CometExecRule { val allExecs: Map[Class[_ <: SparkPlan], CometOperatorSerde[_]] = nativeExecs ++ sinks + /** + * Output path of the write that a `WriteFilesExec` belongs to, copied from the enclosing + * `InsertIntoHadoopFsRelationCommand`. `WriteFilesExec` itself has no output path, so this is + * how `CometWriteFiles` learns the target filesystem - and, by its absence, that a write comes + * from some other V1 write command. + */ + val WRITE_OUTPUT_PATH: TreeNodeTag[String] = TreeNodeTag[String]("comet.writeOutputPath") + /** * Tag set on a `ShuffleExchangeExec` that should be left as a plain Spark shuffle rather than * wrapped in `CometShuffleExchangeExec`. See `tagRedundantColumnarShuffle`. @@ -287,16 +295,11 @@ case class CometExecRule(session: SparkSession) case op if shouldApplySparkToColumnar(conf, op) => convertToComet(op, CometSparkToColumnarExec).getOrElse(op) - // AQE reoptimization looks for `DataWritingCommandExec` or `WriteFilesExec` - // if there is none it would reinsert write nodes, and since Comet remap those nodes - // to Comet counterparties the write nodes are twice to the plan. - // Checking if AQE inserted another write Command on top of existing write command - case _ @DataWritingCommandExec(_, w: WriteFilesExec) - if w.child.isInstanceOf[CometNativeWriteExec] => - w.child - - case op: DataWritingCommandExec => - convertToComet(op, CometDataWritingCommand).getOrElse(op) + // Replace only the per-task write with Comet, leaving DataWritingCommandExec (and therefore + // Spark's commit protocol, stats trackers and SaveMode handling) in place. CometWriteFiles + // declines the write on Spark 3.x, where the node cannot be replaced at all. + case w: WriteFilesExec => + convertToComet(w, CometWriteFiles).getOrElse(w) // For AQE broadcast stage on a Comet broadcast exchange case s @ BroadcastQueryStageExec(_, _: CometBroadcastExchangeExec, _) => @@ -372,14 +375,14 @@ case class CometExecRule(session: SparkSession) op match { case _: CometPlan | _: AQEShuffleReadExec | _: BroadcastExchangeExec | _: BroadcastQueryStageExec | _: AdaptiveSparkPlanExec | _: ExecutedCommandExec | - _: V2CommandExec | _: WriteFilesExec => + _: V2CommandExec | _: WriteFilesExec | _: DataWritingCommandExec => // Some execs should never be replaced. We include // these cases specially here so we do not add a misleading 'info' message. - // WriteFilesExec is always wrapped by DataWritingCommandExec (via Spark's V1Writes - // rule); the parent case converts the whole write to CometNativeWriteExec and - // unwraps WriteFilesExec inside convertToComet. Tagging WriteFilesExec here would - // produce a spurious "WriteFilesExec is not supported" fallback reason (and a warning - // when COMET_EXPLAIN_FALLBACK_LOG_ENABLED=true) even when the write is fully native. + // DataWritingCommandExec is deliberately left in the plan even for a fully native + // write - Comet replaces only its WriteFilesExec child (see CometWriteFiles) - so + // tagging it would report an accelerated write as a fallback. WriteFilesExec that + // reaches this point was already offered to CometWriteFiles by the case above, so it + // has a fallback reason already and must not be tagged again. op case _ => // The operator was not converted to a Comet plan and no serde handler claimed it, so @@ -401,6 +404,17 @@ case class CometExecRule(session: SparkSession) } } + // `WriteFilesExec` does not carry the write's output path, but CometWriteFiles needs it to + // decide whether the target filesystem is supported. Record it from the enclosing command + // before the bottom-up walk reaches the write node. The absence of the tag also tells + // CometWriteFiles that the write is not an InsertIntoHadoopFsRelationCommand and must be + // declined. + plan.foreach { + case DataWritingCommandExec(cmd: InsertIntoHadoopFsRelationCommand, w: WriteFilesExec) => + w.setTagValue(CometExecRule.WRITE_OUTPUT_PATH, cmd.outputPath.toString) + case _ => + } + plan.transformUp { case op => val converted = convertNode(op) // Replace SubqueryBroadcastExec with CometSubqueryBroadcastExec in DPP expressions @@ -681,14 +695,14 @@ case class CometExecRule(session: SparkSession) firstNativeOp = true } - // CometNativeWriteExec is special: it has two separate plans: + // CometWriteFilesExec is special: it has two separate plans: // 1. A protobuf plan (nativeOp) describing the write operation // 2. A Spark plan (child) that produces the data to write // The serializedPlanOpt is a def that always returns Some(...) by serializing // nativeOp on-demand, so it doesn't need convertBlock(). However, its child // (e.g., CometNativeScanExec) may need its own serialization. Reset the flag // so children can start their own native execution blocks. - if (op.isInstanceOf[CometNativeWriteExec]) { + if (op.isInstanceOf[CometWriteFilesExec]) { firstNativeOp = true } @@ -725,12 +739,7 @@ case class CometExecRule(session: SparkSession) // children are CometNativeExec. This prevents runtime failures when the native operator // expects Arrow arrays but receives non-Arrow data (e.g., OnHeapColumnVector). if (serde.requiresNativeChildren && op.children.nonEmpty) { - // Get the actual data-producing children (unwrap WriteFilesExec if present) - val dataProducingChildren = op.children.flatMap { - case writeFiles: WriteFilesExec => Seq(writeFiles.child) - case other => Seq(other) - } - if (!dataProducingChildren.forall(_.isInstanceOf[CometNativeExec])) { + if (!op.children.forall(_.isInstanceOf[CometNativeExec])) { withFallbackReason( op, "Cannot perform native operation because input is not in Arrow format") diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala index ec277cfc7b..7353788474 100644 --- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala +++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala @@ -22,7 +22,7 @@ package org.apache.comet.rules import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.util.sideBySide -import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometPlan, CometSparkToColumnarExec} import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.comet.shims.{MapInBatchInfo, ShimCometMapInBatch} import org.apache.spark.sql.execution.{ColumnarToRowExec, RowToColumnarExec, SparkPlan} @@ -87,10 +87,6 @@ case class EliminateRedundantTransitions(session: SparkSession) // and CometSparkToColumnarExec sparkToColumnar.child } - // Remove unnecessary transition for native writes - // Write should be final operation in the plan - case ColumnarToRowExec(nativeWrite: CometNativeWriteExec) => - nativeWrite case c @ ColumnarToRowExec(child) if hasCometNativeChild(child) => val op = createColumnarToRowExec(child) if (c.logicalLink.isEmpty) { diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala deleted file mode 100644 index 8157f28682..0000000000 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala +++ /dev/null @@ -1,223 +0,0 @@ -/* - * 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.serde.operator - -import java.net.URI -import java.util.Locale - -import scala.jdk.CollectionConverters._ - -import org.apache.parquet.hadoop.ParquetOutputFormat -import org.apache.spark.SparkException -import org.apache.spark.sql.comet.{CometNativeExec, CometNativeWriteExec} -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} -import org.apache.comet.CometSparkSessionExtensions.withFallbackReason -import org.apache.comet.objectstore.NativeConfig -import org.apache.comet.serde.{CometOperatorSerde, Incompatible, OperatorOuterClass, SupportLevel, Unsupported} -import org.apache.comet.serde.OperatorOuterClass.Operator -import org.apache.comet.serde.QueryPlanSerde.serializeDataType - -/** - * CometOperatorSerde implementation for DataWritingCommandExec that converts Parquet write - * operations to use Comet's native Parquet writer. - */ -object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec] { - - private val supportedCompressionCodes = - Set("none", "uncompressed", "snappy", "lz4", "zstd", "gzip") - - override def enabledConfig: Option[ConfigEntry[Boolean]] = - Some(CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED) - - // Native writes require Arrow-formatted input data. If the scan falls back to Spark - // (e.g., due to unsupported complex types), the write must also fall back. - override def requiresNativeChildren: Boolean = true - - override def getSupportLevel(op: DataWritingCommandExec): SupportLevel = { - op.cmd match { - case cmd: InsertIntoHadoopFsRelationCommand => - cmd.fileFormat match { - case _: ParquetFileFormat => - if (!cmd.outputPath.toString.startsWith("file:") && !cmd.outputPath.toString - .startsWith("hdfs:")) { - return Unsupported(Some("Supported output filesystems: local, HDFS")) - } - - if (cmd.bucketSpec.isDefined) { - return Unsupported(Some("Bucketed writes are not supported")) - } - - if (cmd.partitionColumns.nonEmpty || cmd.staticPartitions.nonEmpty) { - return Unsupported(Some("Partitioned writes are not supported")) - } - - val codec = parseCompressionCodec(cmd) - if (!supportedCompressionCodes.contains(codec)) { - return Unsupported(Some(s"Unsupported compression codec: $codec")) - } - - Incompatible(Some("Parquet write support is highly experimental")) - case _ => - Unsupported(Some("Only Parquet writes are supported")) - } - case other => - Unsupported(Some(s"Unsupported write command: ${other.getClass}")) - } - } - - override def convert( - op: DataWritingCommandExec, - builder: Operator.Builder, - childOp: Operator*): Option[OperatorOuterClass.Operator] = { - - try { - val cmd = op.cmd.asInstanceOf[InsertIntoHadoopFsRelationCommand] - - val scanOp = OperatorOuterClass.Scan - .newBuilder() - .setSource(cmd.query.nodeName) - - // Add fields from the query output schema - val scanTypes = cmd.query.output.flatMap { attr => - serializeDataType(attr.dataType) - } - - if (scanTypes.length != cmd.query.output.length) { - withFallbackReason(op, "Cannot serialize data types for native write") - return None - } - - scanTypes.foreach(scanOp.addFields) - - val scanOperator = Operator - .newBuilder() - .setPlanId(op.id) - .setScan(scanOp.build()) - .build() - - val outputPath = cmd.outputPath.toString - - val codec = parseCompressionCodec(cmd) match { - case "snappy" => OperatorOuterClass.CompressionCodec.Snappy - case "lz4" => OperatorOuterClass.CompressionCodec.Lz4 - case "zstd" => OperatorOuterClass.CompressionCodec.Zstd - case "gzip" => OperatorOuterClass.CompressionCodec.Gzip - case "none" | "uncompressed" => OperatorOuterClass.CompressionCodec.None - case other => - withFallbackReason(op, s"Unsupported compression codec: $other") - return None - } - - val writerOpBuilder = OperatorOuterClass.ParquetWriter - .newBuilder() - .setOutputPath(outputPath) - .setCompression(codec) - .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 - - // Collect S3/cloud storage configurations - val session = op.session - val hadoopConf = session.sessionState.newHadoopConfWithOptions(cmd.options) - val objectStoreOptions = - NativeConfig.extractObjectStoreOptions(hadoopConf, URI.create(outputPath)) - objectStoreOptions.foreach { case (key, value) => - writerOpBuilder.putObjectStoreOptions(key, value) - } - - val writerOp = writerOpBuilder.build() - - val writerOperator = Operator - .newBuilder() - .setPlanId(op.id) - .addChildren(scanOperator) - .setParquetWriter(writerOp) - .build() - - Some(writerOperator) - } catch { - case e: Exception => - withFallbackReason( - op, - "Failed to convert DataWritingCommandExec to native execution: " + - s"${e.getMessage}") - None - } - } - - override def createExec(nativeOp: Operator, op: DataWritingCommandExec): CometNativeExec = { - val cmd = op.cmd.asInstanceOf[InsertIntoHadoopFsRelationCommand] - val outputPath = cmd.outputPath.toString - - // Get the child plan from the WriteFilesExec or use the child directly - val childPlan = op.child match { - case writeFiles: WriteFilesExec => - // The WriteFilesExec child should already be a Comet operator - writeFiles.child - case other => - // Fallback: use the child directly - other - } - - // Create FileCommitProtocol for atomic writes - val jobId = java.util.UUID.randomUUID().toString - val committer = - try { - // Use Spark's SQLHadoopMapReduceCommitProtocol - val committerClass = - 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]) - } catch { - case e: Exception => - throw new SparkException(s"Could not instantiate FileCommitProtocol: ${e.getMessage}") - } - - CometNativeWriteExec(nativeOp, childPlan, outputPath, cmd.mode, committer, jobId) - } - - private def parseCompressionCodec(cmd: InsertIntoHadoopFsRelationCommand) = { - // `compression`, `parquet.compression` (i.e., ParquetOutputFormat.COMPRESSION), and - // `spark.sql.parquet.compression.codec` are in order of precedence from highest to - // lowest, matching Spark's own ParquetOptions.compressionCodecClassName. - cmd.options - .get("compression") - .orElse(cmd.options.get(ParquetOutputFormat.COMPRESSION)) - .getOrElse( - SQLConf.get.getConfString( - SQLConf.PARQUET_COMPRESSION.key, - SQLConf.PARQUET_COMPRESSION.defaultValueString)) - .toLowerCase(Locale.ROOT) - } - -} diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometWriteFiles.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometWriteFiles.scala new file mode 100644 index 0000000000..ef1cc335c6 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometWriteFiles.scala @@ -0,0 +1,178 @@ +/* + * 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.serde.operator + +import java.net.URI +import java.util.Locale + +import org.apache.parquet.hadoop.ParquetOutputFormat +import org.apache.spark.sql.comet.{CometNativeExec, CometWriteFilesExec} +import org.apache.spark.sql.execution.datasources.WriteFilesExec +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.{CometConf, ConfigEntry} +import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, withFallbackReason} +import org.apache.comet.objectstore.NativeConfig +import org.apache.comet.rules.CometExecRule +import org.apache.comet.serde.{CometOperatorSerde, Incompatible, OperatorOuterClass, SupportLevel, Unsupported} +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.QueryPlanSerde.serializeDataType + +/** + * Serde for Spark's `WriteFilesExec`, replacing the per-task Parquet write with Comet's native + * writer while leaving the surrounding write framework (commit protocol, stats trackers, SaveMode + * handling, `_SUCCESS`) to Spark. See [[CometWriteFilesExec]] for how the two fit together. + */ +object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { + + private val supportedCompressionCodecs = + Set("none", "uncompressed", "snappy", "lz4", "zstd", "gzip") + + override def enabledConfig: Option[ConfigEntry[Boolean]] = + Some(CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED) + + // Native writes require Arrow-formatted input data. If the query falls back to Spark + // (e.g., due to unsupported complex types), the write must also fall back. + override def requiresNativeChildren: Boolean = true + + override def getSupportLevel(op: WriteFilesExec): SupportLevel = { + // `V1WritesUtils.getWriteFilesOpt` matches the `WriteFilesExecBase` trait on Spark 4.0+, which + // is what makes Spark route the write through CometWriteFilesExec. Spark 3.x matches the + // concrete `WriteFilesExec` case class instead, so a Comet node would be silently ignored and + // the write would fall into FileFormatWriter's non-planned, row-based branch. + if (!isSpark40Plus) { + return Unsupported(Some("Native Parquet writes require Spark 4.0 or later")) + } + + if (!op.fileFormat.isInstanceOf[ParquetFileFormat]) { + return Unsupported(Some("Only Parquet writes are supported")) + } + + // The write node does not carry the output path, so CometExecRule tags it from the enclosing + // InsertIntoHadoopFsRelationCommand. An absent tag means this write belongs to some other V1 + // write command (a Hive insert, for example) whose semantics Comet has not been verified + // against, so decline it. + val outputPath = outputPathOf(op) match { + case Some(path) => path + case None => + return Unsupported(Some("Only InsertIntoHadoopFsRelationCommand writes are supported")) + } + + if (!outputPath.startsWith("file:") && !outputPath.startsWith("hdfs:")) { + return Unsupported(Some("Supported output filesystems: local, HDFS")) + } + + if (op.bucketSpec.isDefined) { + return Unsupported(Some("Bucketed writes are not supported")) + } + + if (op.partitionColumns.nonEmpty || op.staticPartitions.nonEmpty) { + return Unsupported(Some("Partitioned writes are not supported")) + } + + val codec = parseCompressionCodec(op) + if (!supportedCompressionCodecs.contains(codec)) { + return Unsupported(Some(s"Unsupported compression codec: $codec")) + } + + Incompatible(Some("Parquet write support is highly experimental")) + } + + override def convert( + op: WriteFilesExec, + builder: Operator.Builder, + childOp: Operator*): Option[OperatorOuterClass.Operator] = { + + // The native write plan reads from an Arrow stream fed by the already-native child plan, so + // its input is a Scan carrying the child's output schema rather than `childOp`. + val scanOp = OperatorOuterClass.Scan + .newBuilder() + .setSource(op.child.nodeName) + + val scanTypes = op.child.output.flatMap { attr => serializeDataType(attr.dataType) } + if (scanTypes.length != op.child.output.length) { + withFallbackReason(op, "Cannot serialize data types for native write") + return None + } + scanTypes.foreach(scanOp.addFields) + + val scanOperator = Operator + .newBuilder() + .setPlanId(op.id) + .setScan(scanOp.build()) + .build() + + val codec = parseCompressionCodec(op) match { + case "snappy" => OperatorOuterClass.CompressionCodec.Snappy + case "lz4" => OperatorOuterClass.CompressionCodec.Lz4 + case "zstd" => OperatorOuterClass.CompressionCodec.Zstd + case "gzip" => OperatorOuterClass.CompressionCodec.Gzip + case "none" | "uncompressed" => OperatorOuterClass.CompressionCodec.None + case other => + withFallbackReason(op, s"Unsupported compression codec: $other") + return None + } + + // `output_path` and `column_names` are filled in per task by CometWriteFilesExec: the path + // comes from the commit protocol and the names from WriteJobDescription.dataColumns, neither + // of which is known at planning time. + val writerOpBuilder = OperatorOuterClass.ParquetWriter + .newBuilder() + .setCompression(codec) + + // getSupportLevel already declined the write if the tag is absent, so this cannot be empty. + outputPathOf(op).foreach { outputPath => + val hadoopConf = op.session.sessionState.newHadoopConfWithOptions(op.options) + NativeConfig + .extractObjectStoreOptions(hadoopConf, URI.create(outputPath)) + .foreach { case (key, value) => writerOpBuilder.putObjectStoreOptions(key, value) } + } + + Some( + Operator + .newBuilder() + .setPlanId(op.id) + .addChildren(scanOperator) + .setParquetWriter(writerOpBuilder.build()) + .build()) + } + + override def createExec(nativeOp: Operator, op: WriteFilesExec): CometNativeExec = + CometWriteFilesExec(nativeOp, originalPlan = op, child = op.child) + + /** The write's output path, recorded on the node by `CometExecRule`. */ + private def outputPathOf(op: WriteFilesExec): Option[String] = + op.getTagValue(CometExecRule.WRITE_OUTPUT_PATH) + + private def parseCompressionCodec(op: WriteFilesExec): String = { + // `compression`, `parquet.compression` (i.e., ParquetOutputFormat.COMPRESSION), and + // `spark.sql.parquet.compression.codec` are in order of precedence from highest to + // lowest, matching Spark's own ParquetOptions.compressionCodecClassName. + op.options + .get("compression") + .orElse(op.options.get(ParquetOutputFormat.COMPRESSION)) + .getOrElse( + SQLConf.get.getConfString( + SQLConf.PARQUET_COMPRESSION.key, + SQLConf.PARQUET_COMPRESSION.defaultValueString)) + .toLowerCase(Locale.ROOT) + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala index ed94d72869..0f43380b9e 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala @@ -84,26 +84,6 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM } } - /** - * Reports native Parquet writer SQL metrics to Spark's task-level - * [[org.apache.spark.executor.OutputMetrics]] so the Spark UI Stages tab Output column shows - * bytes and records written. - * - * Must be registered on the task thread before [[org.apache.comet.CometExecIterator]] so - * Spark's completion listener stack invokes the iterator `close` (final SQL metric update) - * before this listener runs. - */ - def reportNativeWriteOutputMetrics(ctx: TaskContext): Unit = { - ctx.addTaskCompletionListener[Unit] { _ => - metrics.get("bytes_written").foreach { m => - ctx.taskMetrics().outputMetrics.setBytesWritten(m.value) - } - metrics.get("rows_written").foreach { m => - ctx.taskMetrics().outputMetrics.setRecordsWritten(m.value) - } - } - } - /** * Gets a child node. Called from native. */ diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala deleted file mode 100644 index f0d10b1766..0000000000 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala +++ /dev/null @@ -1,371 +0,0 @@ -/* - * 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.comet - -import scala.jdk.CollectionConverters._ - -import org.apache.hadoop.fs.Path -import org.apache.hadoop.mapreduce.{Job, TaskAttemptContext, TaskAttemptID, TaskID, TaskType} -import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl -import org.apache.spark.TaskContext -import org.apache.spark.internal.io.{FileCommitProtocol, FileNameSpec} -import org.apache.spark.rdd.RDD -import org.apache.spark.sql.SaveMode -import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.comet.execution.arrow.CometArrowStream -import org.apache.spark.sql.comet.util.{Utils => CometUtils} -import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} -import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} -import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} -import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.spark.util.Utils - -import com.google.protobuf.CodedOutputStream - -import org.apache.comet.CometExecIterator -import org.apache.comet.serde.OperatorOuterClass.Operator - -/** - * Comet physical operator for native Parquet write operations with FileCommitProtocol support. - * - * This operator writes data to Parquet files using the native Comet engine. It integrates with - * Spark's FileCommitProtocol to provide atomic writes with proper staging and commit semantics. - * - * The implementation includes support for Spark's file commit protocol through work_dir, job_id, - * and task_attempt_id parameters that can be set in the operator. When work_dir is set, files are - * written to a temporary location that can be atomically committed later. - * - * @param nativeOp - * The native operator representing the write operation (template, will be modified per task) - * @param child - * The child operator providing the data to write - * @param outputPath - * The path where the Parquet file will be written - * @param mode - * The Spark SaveMode governing target-exists behavior (Append / Overwrite / ErrorIfExists / - * Ignore). Comet takes over Spark's DataWritingCommandExec so we must apply these semantics - * here - a direct port of Spark's InsertIntoHadoopFsRelationCommand.run() doInsertion logic. - * @param committer - * FileCommitProtocol for atomic writes. If None, files are written directly. - * @param jobTrackerID - * Unique identifier for this write job - */ -case class CometNativeWriteExec( - nativeOp: Operator, - child: SparkPlan, - outputPath: String, - mode: SaveMode, - committer: Option[FileCommitProtocol] = None, - jobTrackerID: String = Utils.createTempDir().getName) - extends CometNativeExec - with UnaryExecNode { - - override def originalPlan: SparkPlan = child - - // Accumulator to collect TaskCommitMessages from all tasks - // Must be eagerly initialized on driver, not lazy - @transient private val taskCommitMessagesAccum = - sparkContext.collectionAccumulator[FileCommitProtocol.TaskCommitMessage]("taskCommitMessages") - - override def serializedPlanOpt: SerializedPlan = { - val size = nativeOp.getSerializedSize - val bytes = new Array[Byte](size) - val codedOutput = CodedOutputStream.newInstance(bytes) - nativeOp.writeTo(codedOutput) - codedOutput.checkNoSpaceLeft() - SerializedPlan(Some(bytes)) - } - - override def withNewChildInternal(newChild: SparkPlan): SparkPlan = - copy(child = newChild) - - override def nodeName: String = "CometNativeWrite" - - override lazy val metrics: Map[String, SQLMetric] = Map( - "files_written" -> SQLMetrics.createMetric(sparkContext, "number of written data files"), - "bytes_written" -> SQLMetrics.createSizeMetric(sparkContext, "written data"), - "rows_written" -> SQLMetrics.createMetric(sparkContext, "number of written rows")) - - override def doExecute(): RDD[InternalRow] = { - // Setup job if committer is present - committer.foreach { c => - val jobContext = createJobContext() - c.setupJob(jobContext) - } - - // Execute the native write with commit protocol - val resultRDD = doExecuteColumnar() - - // Force execution by consuming all batches - resultRDD - .mapPartitions { iter => - iter.foreach(_.close()) - Iterator.empty - } - .count() - - // Extract write statistics from metrics - val filesWritten = metrics("files_written").value - val bytesWritten = metrics("bytes_written").value - val rowsWritten = metrics("rows_written").value - - // Collect TaskCommitMessages from accumulator - val commitMessages = taskCommitMessagesAccum.value.asScala.toSeq - - // Commit job with collected TaskCommitMessages - committer.foreach { c => - val jobContext = createJobContext() - try { - c.commitJob(jobContext, commitMessages) - logInfo( - s"Successfully committed write job to $outputPath: " + - s"$filesWritten files, $bytesWritten bytes, $rowsWritten rows") - } catch { - case e: Exception => - logError("Failed to commit job, aborting", e) - c.abortJob(jobContext) - throw e - } - } - - // Return empty RDD as write operations don't return data - sparkContext.emptyRDD[InternalRow] - } - - override def doExecuteColumnar(): RDD[ColumnarBatch] = { - // Comet replaces DataWritingCommandExec entirely, so Spark's - // InsertIntoHadoopFsRelationCommand.run() never runs. That method is where Spark handles - // SaveMode semantics (path-exists check, delete-before-Overwrite, Ignore short-circuit) - - // port the non-partitioned, non-catalog branch of that logic here. See Spark 3.5's - // InsertIntoHadoopFsRelationCommand.run doInsertion match. This runs on the driver before - // any executor tasks fire, mirroring where Spark does the delete. - if (!prepareOutputPathForMode()) { - logInfo(s"Skipping insertion into $outputPath - already exists (SaveMode.$mode)") - return sparkContext.emptyRDD[ColumnarBatch] - } - - // Get the input data from the child operator - val childRDD = if (child.supportsColumnar) { - child.executeColumnar() - } else { - // If child doesn't support columnar, convert to columnar - child.execute().mapPartitionsInternal { _ => - // TODO this could delegate to CometRowToColumnar, but maybe Comet - // does not need to support this case? - throw new UnsupportedOperationException( - "Row-based child operators not yet supported for native write") - } - } - - // Capture metadata before the transformation - val numPartitions = childRDD.getNumPartitions - val numOutputCols = child.output.length - val capturedCommitter = committer - val capturedJobTrackerID = jobTrackerID - val capturedNativeOp = nativeOp - val capturedAccumulator = taskCommitMessagesAccum // Capture accumulator for use in tasks - - // Execute native write operation with task-level commit protocol - childRDD.mapPartitionsInternal { iter => - val partitionId = org.apache.spark.TaskContext.getPartitionId() - val taskAttemptId = org.apache.spark.TaskContext.get().taskAttemptId() - - // Setup task-level commit protocol if provided - val (workDir, taskContext, commitMsg) = capturedCommitter - .map { committer => - val taskContext = - createTaskContext(capturedJobTrackerID, partitionId, taskAttemptId.toInt) - - // Setup task - this creates the temporary working directory - committer.setupTask(taskContext) - - // Get the work directory for temp files - // Spark 4.1 made the (taskContext, dir, ext: String) overload throw by default; - // the FileNameSpec overload is the supported one and exists in 3.4+. - val workPath = committer.newTaskTempFile(taskContext, None, FileNameSpec("", "")) - val workDir = new Path(workPath).getParent.toString - - (Some(workDir), Some((committer, taskContext)), null) - } - .getOrElse((None, None, null)) - - // Modify the native operator to include task-specific parameters - val modifiedNativeOp = if (workDir.isDefined) { - val parquetWriter = capturedNativeOp.getParquetWriter.toBuilder - .setWorkDir(workDir.get) - .setJobId(capturedJobTrackerID) - .setTaskAttemptId(taskAttemptId.toInt) - .build() - - capturedNativeOp.toBuilder.setParquetWriter(parquetWriter).build() - } else { - capturedNativeOp - } - - val nativeMetrics = CometMetricNode.fromCometPlan(this) - // Register before CometExecIterator so completion listeners run after iterator close - // (Spark runs task completion callbacks in reverse registration order). - Option(TaskContext.get()).foreach(nativeMetrics.reportNativeWriteOutputMetrics) - - val size = modifiedNativeOp.getSerializedSize - val planBytes = new Array[Byte](size) - val codedOutput = CodedOutputStream.newInstance(planBytes) - modifiedNativeOp.writeTo(codedOutput) - codedOutput.checkNoSpaceLeft() - - val execIterator = new CometExecIterator( - CometExec.newIterId, - CometArrowStream.inputObjects( - iter, - CometUtils.fromAttributes(child.output), - "CometNativeWriteExec"), - numOutputCols, - planBytes, - nativeMetrics, - numPartitions, - partitionId, - None, - Seq.empty) - - // Wrap the iterator to handle task commit/abort and capture TaskCommitMessage - new Iterator[ColumnarBatch] { - private var completed = false - private var thrownException: Option[Throwable] = None - - override def hasNext: Boolean = { - val result = - try { - execIterator.hasNext - } catch { - case e: Throwable => - thrownException = Some(e) - handleTaskEnd() - throw e - } - - if (!result && !completed) { - handleTaskEnd() - } - - result - } - - override def next(): ColumnarBatch = { - try { - execIterator.next() - } catch { - case e: Throwable => - thrownException = Some(e) - handleTaskEnd() - throw e - } - } - - private def handleTaskEnd(): Unit = { - if (!completed) { - completed = true - - // Handle commit or abort based on whether an exception was thrown - taskContext.foreach { case (committer, ctx) => - try { - if (thrownException.isEmpty) { - // Commit the task and add message to accumulator - val message = committer.commitTask(ctx) - capturedAccumulator.add(message) - logDebug(s"Task ${ctx.getTaskAttemptID} committed successfully") - } else { - // Abort the task - committer.abortTask(ctx) - val exMsg = thrownException.get.getMessage - logWarning(s"Task ${ctx.getTaskAttemptID} aborted due to exception: $exMsg") - } - } catch { - case e: Exception => - // Log the commit/abort exception but don't mask the original exception - logError(s"Error during task commit/abort: ${e.getMessage}", e) - if (thrownException.isEmpty) { - // If no original exception, propagate the commit/abort exception - throw e - } - } - } - } - } - } - } - } - - /** Create a JobContext for the write job */ - private def createJobContext(): Job = { - val job = Job.getInstance() - job.setJobID(new org.apache.hadoop.mapreduce.JobID(jobTrackerID, 0)) - job - } - - /** - * Applies SaveMode semantics to the output path before the write starts. Returns `true` when - * the write should proceed and `false` when it should be skipped (Ignore + existing target). - * For ErrorIfExists throws when the target already exists. For Overwrite deletes the target so - * the writer can produce a clean directory. Ported from Spark's - * InsertIntoHadoopFsRelationCommand.run() doInsertion logic, minus the partition/catalog paths - * that Comet does not support. - */ - private def prepareOutputPathForMode(): Boolean = { - val path = new Path(outputPath) - val hadoopConf = sparkContext.hadoopConfiguration - val fs = path.getFileSystem(hadoopConf) - val qualifiedOutputPath = path.makeQualified(fs.getUri, fs.getWorkingDirectory) - - mode match { - case SaveMode.Append => - true - case SaveMode.ErrorIfExists => - if (fs.exists(qualifiedOutputPath)) { - throw QueryCompilationErrors.outputPathAlreadyExistsError(qualifiedOutputPath) - } - true - case SaveMode.Overwrite => - if (fs.exists(qualifiedOutputPath)) { - val deleted = committer match { - case Some(c) => c.deleteWithJob(fs, qualifiedOutputPath, true) - case None => fs.delete(qualifiedOutputPath, true) - } - if (!deleted) { - throw QueryExecutionErrors.cannotClearOutputDirectoryError(qualifiedOutputPath) - } - } - true - case SaveMode.Ignore => - !fs.exists(qualifiedOutputPath) - } - } - - /** Create a TaskAttemptContext for a specific task */ - private def createTaskContext( - jobId: String, - partitionId: Int, - attemptNumber: Int): TaskAttemptContext = { - val job = Job.getInstance() - val taskAttemptID = new TaskAttemptID( - new TaskID(new org.apache.hadoop.mapreduce.JobID(jobId, 0), TaskType.REDUCE, partitionId), - attemptNumber) - new TaskAttemptContextImpl(job.getConfiguration, taskAttemptID) - } -} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala new file mode 100644 index 0000000000..6e899c54f8 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala @@ -0,0 +1,321 @@ +/* + * 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.comet + +import java.util.Date + +import scala.jdk.CollectionConverters._ + +import org.apache.hadoop.mapreduce.{TaskAttemptContext, TaskAttemptID, TaskID, TaskType} +import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl +import org.apache.spark.TaskContext +import org.apache.spark.internal.Logging +import org.apache.spark.internal.io.{FileCommitProtocol, FileNameSpec, SparkHadoopWriterUtils} +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.comet.execution.arrow.CometArrowStream +import org.apache.spark.sql.comet.util.{Utils => CometUtils} +import org.apache.spark.sql.connector.write.WriterCommitMessage +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.{BasicWriteTaskStatsTracker, ExecutedWriteSummary, WriteFilesSpec, WriteJobDescription, WriteTaskResult, WriteTaskStatsTracker} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.Utils + +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.shims.ShimCometWriteFilesExec + +/** + * Comet's replacement for Spark's `WriteFilesExec`: writes Parquet files natively for one task, + * and nothing else. + * + * Everything around the per-task write stays on Spark's side. Because this node extends + * `WriteFilesExecBase` (see [[ShimCometWriteFilesExec]]), `V1WritesUtils.getWriteFilesOpt` finds + * it, so `InsertIntoHadoopFsRelationCommand.run` and `FileFormatWriter` continue to own: + * + * - SaveMode semantics and the delete-before-overwrite of the target directory + * - committer instantiation (including `spark.sql.sources.commitProtocolClass`), `setupJob`, + * `commitJob`/`abortJob`, `onTaskCommit`, and the `_SUCCESS` marker + * - dynamic partition overwrite and custom partition locations + * - `WriteJobStatsTracker` aggregation, SQL metrics, and catalog statistics/cache refresh + * + * This node mirrors `FileFormatWriter.executeTask` for the parts Comet must do itself: build the + * `TaskAttemptContext`, ask the commit protocol where to write, run the native writer, drive the + * stats trackers, and commit or abort the task. Notably the path handed back by + * `FileCommitProtocol.newTaskTempFile` is used verbatim, so staging directories, task-attempt + * isolation under speculation, and committers that track individual files all behave as they do + * for Spark's own writer. + * + * @param nativeOp + * Template for the native write plan. `output_path` is a placeholder here and is replaced per + * task with the path the commit protocol chose. + * @param originalPlan + * The `WriteFilesExec` this node replaced. This must be the write node rather than the data + * subtree: `CometExecRule` copies `originalPlan`'s logical link onto every `CometExec`, and + * pointing it at the child would link the write node to the child's logical plan, which makes + * AQE mistake it for the child's query stage and re-wrap it in a second `WriteFilesExec`. + * @param child + * The Comet native operator producing the batches to write. + */ +case class CometWriteFilesExec( + nativeOp: Operator, + override val originalPlan: SparkPlan, + child: SparkPlan) + extends CometNativeExec + with ShimCometWriteFilesExec { + + override def nodeName: String = "CometWriteFiles" + + override lazy val metrics: Map[String, SQLMetric] = Map( + "files_written" -> SQLMetrics.createMetric(sparkContext, "number of written data files"), + "bytes_written" -> SQLMetrics.createSizeMetric(sparkContext, "written data"), + "rows_written" -> SQLMetrics.createMetric(sparkContext, "number of written rows")) + + override def serializedPlanOpt: SerializedPlan = + SerializedPlan(Some(CometExec.serializeNativePlan(nativeOp))) + + override def withNewChildInternal(newChild: SparkPlan): SparkPlan = copy(child = newChild) + + /** + * Spark drives this node through `executeWrite`, never `execute`. `WriteFilesExecBase` already + * throws for `doExecute`, but `CometExec` widens it to a public member that returns a + * `ColumnarToRowExec` result, so the conflict has to be resolved explicitly here. + */ + override def doExecute(): RDD[InternalRow] = + throw new UnsupportedOperationException(s"$nodeName does not support doExecute") + + override protected def doExecuteWrite( + writeFilesSpec: WriteFilesSpec): RDD[WriterCommitMessage] = { + val description = writeFilesSpec.description + val committer = writeFilesSpec.committer + // Same identifier scheme as FileFormatWriter, so committers that parse the job ID agree. + val jobTrackerID = SparkHadoopWriterUtils.createJobTrackerID(new Date()) + + val childRDD = child.executeColumnar() + + // Everything the write task needs is resolved here on the driver and captured by value. The + // closure below must not touch `this`: a CometWriteFilesExec holds `nativeOp` plus the whole + // converted child subtree, each node of which carries its own non-transient protobuf, so + // capturing it would ship a redundant copy of the plan to every executor. Spark's own + // WriteFilesExec.doExecuteWrite avoids this the same way, by delegating to a static + // FileFormatWriter.executeTask. + val taskWrite = NativeWriteTask( + nativeOp = nativeOp, + // Column names come from the write job description, not from the query output: for + // `INSERT INTO t SELECT ...` the query may name columns after the expressions that produced + // them, while the file must carry the target table's column names. + dataColumnNames = description.dataColumns.map(_.name), + childSchema = CometUtils.fromAttributes(child.output), + numPartitions = childRDD.getNumPartitions, + nativeMetrics = CometMetricNode.fromCometPlan(this), + nodeName = nodeName) + + assert( + taskWrite.dataColumnNames.length == child.output.length, + s"Expected ${taskWrite.dataColumnNames.length} data columns to write but the child " + + s"produces ${child.output.length}") + + childRDD.mapPartitionsInternal { batches => + CometWriteFilesExec.executeTask(description, committer, jobTrackerID, taskWrite, batches) + } + } +} + +/** + * The per-task state that [[CometWriteFilesExec.executeTask]] needs, resolved on the driver. + * + * A plain container rather than a closure over the exec node: it copies only these fields, so the + * enclosing plan tree is not kept alive for the task's lifetime or shipped in the task binary. + */ +private[comet] case class NativeWriteTask( + nativeOp: Operator, + dataColumnNames: Seq[String], + childSchema: StructType, + numPartitions: Int, + nativeMetrics: CometMetricNode, + nodeName: String) + +object CometWriteFilesExec extends Logging { + + /** + * Write one task's batches natively and commit or abort it, mirroring the structure of + * `FileFormatWriter.executeTask`. + */ + private[comet] def executeTask( + description: WriteJobDescription, + committer: FileCommitProtocol, + jobTrackerID: String, + taskWrite: NativeWriteTask, + batches: Iterator[ColumnarBatch]): Iterator[WriterCommitMessage] = { + val taskCtx = TaskContext.get() + val sparkPartitionId = taskCtx.partitionId() + val taskAttemptContext = createTaskAttemptContext( + description, + jobTrackerID, + taskCtx.stageId(), + sparkPartitionId, + // Truncation to Int matches FileFormatWriter: the masked low bits are what the Hadoop + // TaskAttemptID accepts, and uniqueness within a job is preserved by the task ID. + taskCtx.taskAttemptId().toInt & Integer.MAX_VALUE) + + committer.setupTask(taskAttemptContext) + val statsTrackers = description.statsTrackers.map(_.newTaskInstance()) + + try { + // Mirrors FileFormatWriter's EmptyDirectoryDataWriter case: an empty input still writes one + // file from partition 0 so that the output carries the schema, but every other empty + // partition produces no file at all. + val writtenFile = if (sparkPartitionId == 0 || batches.hasNext) { + val ext = description.outputWriterFactory.getFileExtension(taskAttemptContext) + // FileNameSpec's "-c000" suffix reproduces Spark's part---c000..parquet + // naming. The file counter is always 0 until file rolling is supported. + val filePath = + committer.newTaskTempFile(taskAttemptContext, None, FileNameSpec("", "-c000" + ext)) + + statsTrackers.foreach(_.newFile(filePath)) + val rowsWritten = writeNatively(taskWrite, filePath, batches, sparkPartitionId) + recordRows(statsTrackers, filePath, rowsWritten) + statsTrackers.foreach(_.closeFile(filePath)) + filePath + } else { + // Drain so the child's native execution completes and releases its resources. + batches.foreach(_.close()) + "no file" + } + + val (taskCommitMessage, taskCommitTime) = Utils.timeTakenMs { + committer.commitTask(taskAttemptContext) + } + logDebug(s"Task ${taskAttemptContext.getTaskAttemptID} committed $writtenFile") + + Iterator( + WriteTaskResult( + taskCommitMessage, + ExecutedWriteSummary( + // Only non-partitioned writes are supported so far, so no partition paths were + // added. Populating this is part of adding partitioned write support. + updatedPartitions = Set.empty, + stats = statsTrackers.map(_.getFinalStats(taskCommitTime))))) + } catch { + case t: Throwable => + Utils.tryLogNonFatalError(committer.abortTask(taskAttemptContext)) + logError(s"Task ${taskAttemptContext.getTaskAttemptID} aborted: ${t.getMessage}", t) + throw t + } + } + + /** + * Run the native write plan for one task, returning the number of rows written. + * + * The row count is taken on the JVM side as batches are pulled into native code. That is exact + * because the native writer consumes its whole input before completing. + */ + private def writeNatively( + taskWrite: NativeWriteTask, + filePath: String, + batches: Iterator[ColumnarBatch], + partitionId: Int): Long = { + val parquetWriter = taskWrite.nativeOp.getParquetWriter.toBuilder + .setOutputPath(filePath) + .clearColumnNames() + .addAllColumnNames(taskWrite.dataColumnNames.asJava) + .build() + val taskOp = taskWrite.nativeOp.toBuilder.setParquetWriter(parquetWriter).build() + + var rowsWritten = 0L + val countingBatches = + CometArrowStream.countingIterator[ColumnarBatch](batches, b => rowsWritten += b.numRows()) + + val execIterator = CometExec.getCometIterator( + CometArrowStream.inputObjects(countingBatches, taskWrite.childSchema, taskWrite.nodeName), + taskWrite.dataColumnNames.length, + taskOp, + taskWrite.nativeMetrics, + taskWrite.numPartitions, + partitionId, + broadcastedHadoopConfForEncryption = None, + encryptedFilePaths = Seq.empty) + + try { + // The native writer emits no batches; draining performs the write. + while (execIterator.hasNext) { + execIterator.next().close() + } + } finally { + execIterator.close() + } + + rowsWritten + } + + /** + * Report `count` rows to each stats tracker. + * + * `WriteTaskStatsTracker.newRow` is a per-row callback, but the only implementation Spark + * ships, [[BasicWriteTaskStatsTracker]], ignores the row and just counts. Comet has columnar + * batches rather than `InternalRow`s here, so it passes an empty row instead of materializing + * every row just to hand it straight back. A tracker that actually inspects row contents would + * therefore see empty rows, so warn rather than silently report wrong statistics. + * + * The loop is per-tracker on the outside so the hot inner loop has a single receiver and no + * per-row closure; the trackers are independent per-file counters, so their relative + * interleaving carries no meaning. + */ + private def recordRows( + statsTrackers: Seq[WriteTaskStatsTracker], + filePath: String, + count: Long): Unit = { + statsTrackers.foreach { tracker => + if (!tracker.isInstanceOf[BasicWriteTaskStatsTracker]) { + logWarning( + s"${tracker.getClass.getName} receives row counts but not row contents from Comet's " + + "native Parquet writer. Set spark.comet.parquet.write.enabled=false if this tracker " + + "needs to inspect written rows.") + } + var i = 0L + while (i < count) { + tracker.newRow(filePath, InternalRow.empty) + i += 1 + } + } + } + + /** Build the `TaskAttemptContext` exactly as `FileFormatWriter.executeTask` does. */ + private def createTaskAttemptContext( + description: WriteJobDescription, + jobTrackerID: String, + sparkStageId: Int, + sparkPartitionId: Int, + sparkAttemptNumber: Int): TaskAttemptContext = { + val jobId = SparkHadoopWriterUtils.createJobID(jobTrackerID, sparkStageId) + val taskId = new TaskID(jobId, TaskType.MAP, sparkPartitionId) + val taskAttemptId = new TaskAttemptID(taskId, sparkAttemptNumber) + + val hadoopConf = description.serializableHadoopConf.value + hadoopConf.set("mapreduce.job.id", jobId.toString) + hadoopConf.set("mapreduce.task.id", taskAttemptId.getTaskID.toString) + hadoopConf.set("mapreduce.task.attempt.id", taskAttemptId.toString) + hadoopConf.setBoolean("mapreduce.task.ismap", true) + hadoopConf.setInt("mapreduce.task.partition", 0) + + new TaskAttemptContextImpl(hadoopConf, taskAttemptId) + } +} diff --git a/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala new file mode 100644 index 0000000000..f6551c60eb --- /dev/null +++ b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala @@ -0,0 +1,36 @@ +/* + * 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.shims + +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.execution.UnaryExecNode + +/** + * Base type for [[org.apache.spark.sql.comet.CometWriteFilesExec]] on Spark 3.x. + * + * Spark 3.x has no `WriteFilesExecBase` trait (added in 4.0): `V1WritesUtils.getWriteFilesOpt` + * matches the concrete `WriteFilesExec` case class, so a Comet node can never be picked up as the + * write node there. Native writes are gated to Spark 4.0+ in `CometExecRule` and this shim exists + * only so that the shared sources compile against 3.x. It mirrors the members that the 4.x + * `WriteFilesExecBase` supplies. + */ +trait ShimCometWriteFilesExec extends UnaryExecNode { + override def output: Seq[Attribute] = Seq.empty +} diff --git a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala new file mode 100644 index 0000000000..028df2f89d --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala @@ -0,0 +1,37 @@ +/* + * 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.shims + +import org.apache.spark.sql.execution.datasources.WriteFilesExecBase + +/** + * Base type for [[org.apache.spark.sql.comet.CometWriteFilesExec]]. + * + * Spark 4.0 factored `WriteFilesExec`'s contract out into the `WriteFilesExecBase` trait, and + * `V1WritesUtils.getWriteFilesOpt` matches on that trait. Extending it is therefore what makes + * Spark recognize Comet's node as the write node and drive it through + * `FileFormatWriter.executeWrite` -> `SparkPlan.executeWrite` -> `doExecuteWrite`, keeping the + * commit protocol, stats trackers and `_SUCCESS` handling on Spark's side. + * + * Spark 3.x has no such trait - `getWriteFilesOpt` matches the concrete `WriteFilesExec` case + * class - so native writes are gated to Spark 4.0+ in `CometExecRule`. The 3.x variant of this + * shim exists only to keep the shared sources compiling. + */ +trait ShimCometWriteFilesExec extends WriteFilesExecBase diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index a1ae1af1d1..260199652c 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -24,25 +24,41 @@ import java.io.File import scala.jdk.CollectionConverters._ import scala.util.{Random, Using} +import org.scalactic.source.Position +import org.scalatest.Tag + import org.apache.hadoop.fs.{FileSystem, Path} import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.metadata.CompressionCodecName import org.apache.parquet.hadoop.util.HadoopInputFile import org.apache.spark.sql.{AnalysisException, CometTestBase, DataFrame, Row, SaveMode} -import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometNativeWriteExec, CometScanExec} +import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometScanExec, CometWriteFilesExec} import org.apache.spark.sql.execution.{FileSourceScanExec, QueryExecution, SparkPlan} import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.WriteFilesExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType -import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus +import org.apache.comet.{CometConf, CometExplainInfo} +import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, SchemaGenOptions} class CometParquetWriterSuite extends CometTestBase { import testImplicits._ + /** + * Native Parquet writes hook into Spark's write path through `WriteFilesExecBase`, which only + * exists in Spark 4.0+. See `CometWriteFilesExec` and the gate in `CometExecRule`. + */ + override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit + pos: Position): Unit = { + super.test(testName, testTags: _*) { + assume(isSpark40Plus, "Comet native Parquet writes require Spark 4.0+") + testFun + } + } + test("partitioned write with empty string partition value") { withTempPath { path => Seq(("", 1), ("a", 2)) @@ -75,10 +91,10 @@ class CometParquetWriterSuite extends CometTestBase { withSQLConf( CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax", - CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + CometConf.COMET_OPERATOR_WRITE_FILES_ALLOW_INCOMPAT.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true") { - writeWithCometNativeWriteExec(inputPath, outputPath) + writeWithCometWriteFilesExec(inputPath, outputPath) verifyWrittenFile(outputPath) } @@ -97,10 +113,10 @@ class CometParquetWriterSuite extends CometTestBase { withSQLConf( CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax", - CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + CometConf.COMET_OPERATOR_WRITE_FILES_ALLOW_INCOMPAT.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true") { - val capturedPlan = writeWithCometNativeWriteExec(inputPath, outputPath) + val capturedPlan = writeWithCometWriteFilesExec(inputPath, outputPath) capturedPlan.foreach { plan => val hasNativeScan = plan.exists { case _: CometNativeScanExec => true @@ -131,11 +147,10 @@ class CometParquetWriterSuite extends CometTestBase { CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", "spark.sql.adaptive.enabled" -> adaptive.toString, SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax", - CometConf.getOperatorAllowIncompatConfigKey( - classOf[DataWritingCommandExec]) -> "true", + CometConf.getOperatorAllowIncompatConfigKey(classOf[WriteFilesExec]) -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true") { - writeWithCometNativeWriteExec(inputPath, outputPath, Some(10)) + writeWithCometWriteFilesExec(inputPath, outputPath, Some(10)) verifyWrittenFile(outputPath) } }) @@ -151,12 +166,12 @@ class CometParquetWriterSuite extends CometTestBase { withSQLConf( CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", - CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true", + CometConf.getOperatorAllowIncompatConfigKey(classOf[WriteFilesExec]) -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", SQLConf.PARQUET_COMPRESSION.key -> codec) { val plan = captureWritePlan(path => df.write.parquet(path), outputPath) - assertHasCometNativeWriteExec(plan) + assertHasCometWriteFilesExec(plan) } checkAnswer(spark.read.parquet(outputPath), df.collect()) @@ -172,14 +187,14 @@ class CometParquetWriterSuite extends CometTestBase { withSQLConf( CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", - CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true", + CometConf.getOperatorAllowIncompatConfigKey(classOf[WriteFilesExec]) -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", SQLConf.PARQUET_COMPRESSION.key -> "snappy") { val plan = captureWritePlan( path => df.write.option("parquet.compression", "gzip").parquet(path), outputPath) - assertHasCometNativeWriteExec(plan) + assertHasCometWriteFilesExec(plan) } checkAnswer(spark.read.parquet(outputPath), df.collect()) @@ -197,7 +212,7 @@ class CometParquetWriterSuite extends CometTestBase { withSQLConf( CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", - CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true", + CometConf.getOperatorAllowIncompatConfigKey(classOf[WriteFilesExec]) -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", SQLConf.PARQUET_COMPRESSION.key -> "zstd") { @@ -208,7 +223,7 @@ class CometParquetWriterSuite extends CometTestBase { .option("parquet.compression", "snappy") .parquet(path), outputPath) - assertHasCometNativeWriteExec(plan) + assertHasCometWriteFilesExec(plan) } checkAnswer(spark.read.parquet(outputPath), df.collect()) @@ -224,12 +239,12 @@ class CometParquetWriterSuite extends CometTestBase { withSQLConf( CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", - CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true", + CometConf.getOperatorAllowIncompatConfigKey(classOf[WriteFilesExec]) -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", SQLConf.PARQUET_COMPRESSION.key -> "lz4_raw") { val plan = captureWritePlan(path => df.write.parquet(path), outputPath) - assertNoCometNativeWriteExec(plan) + assertNoCometWriteFilesExec(plan) } checkAnswer(spark.read.parquet(outputPath), df.collect()) @@ -457,7 +472,7 @@ class CometParquetWriterSuite extends CometTestBase { sourcePath) val plan = captureWritePlan(p => df.write.mode(SaveMode.ErrorIfExists).parquet(p), outputPath) - assertHasCometNativeWriteExec(plan) + assertHasCometWriteFilesExec(plan) checkAnswer(spark.read.parquet(outputPath), df) } } @@ -497,7 +512,7 @@ class CometParquetWriterSuite extends CometTestBase { (1 to 100).map(i => (i, s"str_$i")).toDF("id", "name"), sourcePath) val plan = captureWritePlan(p => df.write.mode(SaveMode.Overwrite).parquet(p), outputPath) - assertHasCometNativeWriteExec(plan) + assertHasCometWriteFilesExec(plan) checkAnswer(spark.read.parquet(outputPath), df) } } @@ -518,7 +533,7 @@ class CometParquetWriterSuite extends CometTestBase { sourcePath) val plan = captureWritePlan(p => replacement.write.mode(SaveMode.Overwrite).parquet(p), outputPath) - assertHasCometNativeWriteExec(plan) + assertHasCometWriteFilesExec(plan) checkAnswer(spark.read.parquet(outputPath), replacement) } } @@ -553,7 +568,7 @@ class CometParquetWriterSuite extends CometTestBase { (1 to 100).map(i => (i, s"str_$i")).toDF("id", "name"), sourcePath) val plan = captureWritePlan(p => df.write.mode(SaveMode.Append).parquet(p), outputPath) - assertHasCometNativeWriteExec(plan) + assertHasCometWriteFilesExec(plan) checkAnswer(spark.read.parquet(outputPath), df) } } @@ -576,7 +591,7 @@ class CometParquetWriterSuite extends CometTestBase { sourcePath) val plan = captureWritePlan(p => second.write.mode(SaveMode.Append).parquet(p), outputPath) - assertHasCometNativeWriteExec(plan) + assertHasCometWriteFilesExec(plan) val filesAfter = listPartFileNames(outputPath) assert( @@ -623,7 +638,7 @@ class CometParquetWriterSuite extends CometTestBase { (1 to 100).map(i => (i, s"str_$i")).toDF("id", "name"), sourcePath) val plan = captureWritePlan(p => df.write.mode(SaveMode.Ignore).parquet(p), outputPath) - assertHasCometNativeWriteExec(plan) + assertHasCometWriteFilesExec(plan) checkAnswer(spark.read.parquet(outputPath), df) } } @@ -651,6 +666,76 @@ class CometParquetWriterSuite extends CometTestBase { } } + test("write creates a _SUCCESS marker") { + // https://github.com/apache/datafusion-comet/issues/2985 - the marker comes from + // HadoopMapReduceCommitProtocol.commitJob, which only runs because Comet leaves + // InsertIntoHadoopFsRelationCommand in the plan. + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + withTempPath { srcDir => + val df = materializeAsCometSource( + (1 to 100).map(i => (i, s"n_$i")).toDF("id", "name"), + new File(srcDir, "src.parquet").getAbsolutePath) + withNativeWriter { + val plan = captureWritePlan(p => df.write.parquet(p), outputPath) + assertHasCometWriteFilesExec(plan) + } + } + assert( + new File(outputPath, "_SUCCESS").exists(), + s"Expected a _SUCCESS marker in $outputPath, found: " + + new File(outputPath).list().mkString(", ")) + } + } + + test("written file names follow Spark's naming convention") { + // The file name comes from FileCommitProtocol.newTaskTempFile and must be used verbatim: + // part---c..parquet. Committers that track individual files + // and tools that parse these names depend on it. + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + withTempPath { srcDir => + val df = materializeAsCometSource( + (1 to 100).map(i => (i, s"n_$i")).toDF("id", "name"), + new File(srcDir, "src.parquet").getAbsolutePath) + withNativeWriter { + withSQLConf(SQLConf.PARQUET_COMPRESSION.key -> "snappy") { + val plan = captureWritePlan(p => df.write.parquet(p), outputPath) + assertHasCometWriteFilesExec(plan) + } + } + } + + val partFiles = listPartFileNames(outputPath) + assert(partFiles.nonEmpty, s"No part files written to $outputPath") + val namePattern = + """part-\d{5}-[0-9a-f\-]{36}-c\d{3}\.snappy\.parquet""".r + partFiles.foreach { name => + assert( + namePattern.pattern.matcher(name).matches(), + s"File name '$name' does not match Spark's part-file naming convention") + } + } + } + + test("INSERT INTO ... SELECT is visible to subsequent reads") { + // https://github.com/apache/datafusion-comet/issues/3521 - reads returned no rows because the + // bespoke write path never refreshed the catalog cache. Spark's command does that itself. + withTable("comet_write_target", "comet_write_source") { + withNativeWriter { + sql("CREATE TABLE comet_write_source(id bigint, name string) USING parquet") + sql("CREATE TABLE comet_write_target(id bigint, name string) USING parquet") + } + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sql("INSERT INTO comet_write_source VALUES (1, 'a'), (2, 'b')") + } + withNativeWriter { + sql("INSERT INTO comet_write_target SELECT id, name FROM comet_write_source") + } + checkAnswer(spark.table("comet_write_target"), Row(1L, "a") :: Row(2L, "b") :: Nil) + } + } + private def createTestData(inputDir: File): String = { val inputPath = new File(inputDir, "input.parquet").getAbsolutePath val schema = FuzzDataGenerator.generateSchema( @@ -672,7 +757,7 @@ class CometParquetWriterSuite extends CometTestBase { private def withNativeWriter(f: => Unit): Unit = { withSQLConf( CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", - CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + CometConf.COMET_OPERATOR_WRITE_FILES_ALLOW_INCOMPAT.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax")(f) } @@ -751,42 +836,45 @@ class CometParquetWriterSuite extends CometTestBase { } } - private def assertHasCometNativeWriteExec(plan: SparkPlan): Unit = { + // CometWriteFilesExec replaces only WriteFilesExec, so it appears in the plan beneath Spark's + // DataWritingCommandExec. A single `plan.foreach` over the whole tree therefore sees it exactly + // once; there is no separate command node to inspect. + private def assertHasCometWriteFilesExec(plan: SparkPlan): Unit = { var nativeWriteCount = 0 plan.foreach { - case _: CometNativeWriteExec => - nativeWriteCount += 1 - case d: DataWritingCommandExec => - d.child.foreach { - case _: CometNativeWriteExec => - nativeWriteCount += 1 - case _ => - } + case _: CometWriteFilesExec => nativeWriteCount += 1 case _ => } assert( nativeWriteCount == 1, - s"Expected exactly one CometNativeWriteExec in the plan, but found $nativeWriteCount:\n${plan.treeString}") + s"Expected exactly one CometWriteFilesExec in the plan, but found $nativeWriteCount:\n${plan.treeString}") + + // The command is left in the plan on purpose for a fully native write, so it must not be + // reported as a fallback - otherwise extended explain tells users an accelerated write was + // not accelerated, and skews the "Comet accelerated N of M operators" count. + plan.foreach { + case d: DataWritingCommandExec => + val reasons = d.getTagValue(CometExplainInfo.FALLBACK_REASONS).getOrElse(Set.empty) + assert( + reasons.isEmpty, + s"A fully native write must not tag ${d.nodeName} as a fallback, got: $reasons") + case _ => + } } - private def assertNoCometNativeWriteExec(plan: SparkPlan): Unit = { + private def assertNoCometWriteFilesExec(plan: SparkPlan): Unit = { val hasNativeWrite = plan.exists { - case _: CometNativeWriteExec => true - case d: DataWritingCommandExec => - d.child.exists { - case _: CometNativeWriteExec => true - case _ => false - } + case _: CometWriteFilesExec => true case _ => false } assert( !hasNativeWrite, - s"Expected no CometNativeWriteExec in the plan, but found one:\n${plan.treeString}") + s"Expected no CometWriteFilesExec in the plan, but found one:\n${plan.treeString}") } - private def writeWithCometNativeWriteExec( + private def writeWithCometWriteFilesExec( inputPath: String, outputPath: String, num_partitions: Option[Int] = None): Option[SparkPlan] = { @@ -796,7 +884,7 @@ class CometParquetWriterSuite extends CometTestBase { path => num_partitions.fold(df)(n => df.repartition(n)).write.parquet(path), outputPath) - assertHasCometNativeWriteExec(plan) + assertHasCometWriteFilesExec(plan) Some(plan) } @@ -871,7 +959,7 @@ class CometParquetWriterSuite extends CometTestBase { withSQLConf( CometConf.COMET_EXEC_ENABLED.key -> "true", // enable experimental native writes - CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + CometConf.COMET_OPERATOR_WRITE_FILES_ALLOW_INCOMPAT.key -> "true", CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", // Disable unsigned small int safety check for ShortType columns CometConf.COMET_PARQUET_UNSIGNED_SMALL_INT_CHECK.key -> "false", @@ -880,9 +968,9 @@ class CometParquetWriterSuite extends CometTestBase { val parquetDf = spark.read.parquet(inputPath) - // Capture plan and verify CometNativeWriteExec is used + // Capture plan and verify CometWriteFilesExec is used val plan = captureWritePlan(path => parquetDf.write.parquet(path), outputPath) - assertHasCometNativeWriteExec(plan) + assertHasCometWriteFilesExec(plan) } // Verify round-trip: read with Spark and Comet, compare results diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala index fdcca8d351..e370e3e006 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala @@ -34,7 +34,7 @@ import org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffle import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.WriteFilesExec import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf @@ -202,8 +202,7 @@ class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper { withSQLConf( CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", - CometConf.getOperatorAllowIncompatConfigKey( - classOf[DataWritingCommandExec]) -> "true", + CometConf.getOperatorAllowIncompatConfigKey(classOf[WriteFilesExec]) -> "true", SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax") { spark.sparkContext.setJobGroup(jobGroupId, "native parquet write output metrics") try {