From 0f946b5bffcd841fb51bdd348d1b5d2421d65b09 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 7 Aug 2026 08:48:11 -0600 Subject: [PATCH 1/2] refactor: hook native Parquet writes into Spark's WriteFilesExec seam Native Parquet writes previously replaced the whole DataWritingCommandExec, so InsertIntoHadoopFsRelationCommand.run never ran and Comet had to re-implement the write framework itself: a hardcoded SQLHadoopMapReduceCommitProtocol, a hand-ported copy of the SaveMode logic, its own commit-message accumulator, and its own commitJob call. Most of the open native-writer bugs are symptoms of that one decision rather than independent defects. Replace only WriteFilesExec instead. On Spark 4.0+ V1WritesUtils.getWriteFilesOpt matches the WriteFilesExecBase trait, so extending it is what makes Spark route the write through CometWriteFilesExec.doExecuteWrite while keeping ownership of the commit protocol, staging, _SUCCESS, dynamic partition overwrite, stats tracker aggregation, and catalog refresh. Comet is left owning just the per-task native encode. Gated to Spark 4.0+: Spark 3.x matches the concrete WriteFilesExec case class, so a Comet node would be silently ignored and the write would drop into FileFormatWriter's non-planned, row-based branch. Other behaviour changes that fall out of the new seam: - File paths come from FileCommitProtocol.newTaskTempFile and are used verbatim, so names match Spark's part---c000..parquet and committers that track individual files work. Previously only the parent directory was kept and the native writer invented its own names. - Column names come from WriteJobDescription.dataColumns rather than the query output, so INSERT INTO t SELECT a+1 writes the target column name. - Byte and row counts come from BasicWriteTaskStatsTracker, which stats files through the FileSystem API and so is correct for HDFS. The native writer's std::fs::metadata call reported 0 there. CometMetricNode's now-redundant reportNativeWriteOutputMetrics is removed. - AQE re-plans the write command's child and re-inserts a WriteFilesExec above the already-converted node; collapse it rather than nesting two native writes, which would write the data twice. Deletes CometNativeWriteExec and CometDataWritingCommand (594 lines). The operator opt-in moves to spark.comet.operator.WriteFilesExec.allowIncompatible, with the previous DataWritingCommandExec key kept as an alternative. --- docs/source/user-guide/latest/installation.md | 4 +- docs/source/user-guide/latest/operators.md | 6 +- .../src/execution/operators/parquet_writer.rs | 46 +-- native/core/src/execution/planner.rs | 8 - native/proto/src/proto/operator.proto | 15 +- .../scala/org/apache/comet/CometConf.scala | 22 +- .../apache/comet/rules/CometExecRule.scala | 68 ++-- .../rules/EliminateRedundantTransitions.scala | 11 +- .../operator/CometDataWritingCommand.scala | 223 ----------- .../serde/operator/CometWriteFiles.scala | 164 ++++++++ .../spark/sql/comet/CometMetricNode.scala | 20 - .../sql/comet/CometNativeWriteExec.scala | 371 ------------------ .../spark/sql/comet/CometWriteFilesExec.scala | 298 ++++++++++++++ .../comet/shims/ShimCometWriteFilesExec.scala | 36 ++ .../comet/shims/ShimCometWriteFilesExec.scala | 37 ++ .../parquet/CometParquetWriterSuite.scala | 171 +++++--- .../sql/comet/CometTaskMetricsSuite.scala | 5 +- 17 files changed, 748 insertions(+), 757 deletions(-) delete mode 100644 spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala create mode 100644 spark/src/main/scala/org/apache/comet/serde/operator/CometWriteFiles.scala delete mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala create mode 100644 spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala create mode 100644 spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala 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..6f1b33b40d 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", + Seq(getOperatorAllowIncompatConfigKey("DataWritingCommandExec"))) /** Create a config to enable a specific operator */ private def createExecEnabledConfig( @@ -896,15 +901,20 @@ 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, + alternatives: Seq[String] = Seq.empty): 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) + val withAlts = alternatives match { + case Nil => builder + case alt +: more => builder.withAlternative(alt, more: _*) + } + withAlts.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..949c9a5fe9 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,28 @@ 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] => + // AQE re-plans the write command's child, which re-inserts a WriteFilesExec above the node + // Comet already converted on a previous pass. Collapse it instead of wrapping a second + // native write around the first: nested writes would write the data twice, and the inner + // node's empty output would be mistaken for a zero-column schema. + case w: WriteFilesExec if w.child.isInstanceOf[CometWriteFilesExec] => 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. Gated to Spark + // 4.0+: `V1WritesUtils.getWriteFilesOpt` matches the `WriteFilesExecBase` trait there, which + // is what makes Spark route the write through this node. 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. + case w: WriteFilesExec if isSpark40Plus => + convertToComet(w, CometWriteFiles).getOrElse(w) + + // On Spark 3.x the write node cannot be replaced at all, so explain why rather than leaving + // a user who opted into native writes with no fallback reason. Only reported when they asked + // for the feature; otherwise WriteFilesExec falls through to the never-replaced list below. + case w: WriteFilesExec if CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.get(conf) => + withFallbackReason(w, "Native Parquet writes require Spark 4.0 or later") + w // For AQE broadcast stage on a Comet broadcast exchange case s @ BroadcastQueryStageExec(_, _: CometBroadcastExchangeExec, _) => @@ -375,11 +395,9 @@ case class CometExecRule(session: SparkSession) _: V2CommandExec | _: WriteFilesExec => // 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. + // A WriteFilesExec that reaches this point was already offered to CometWriteFiles by + // the case above (or skipped on Spark 3.x, where the node cannot be replaced at all), + // 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 +419,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 +710,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 +754,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..ac2b4252cd 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, CometWriteFilesExec} 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,9 +87,12 @@ 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) => + // Remove unnecessary transition for native writes. CometWriteFilesExec reports + // supportsColumnar (it is a CometExec) so that its child stays columnar, but Spark's + // transition rules then wrap it for the row-based DataWritingCommandExec above. That + // wrapper has to go: Spark drives the write through executeWrite on this exact node, and + // ColumnarToRowExec does not implement doExecuteWrite. + case ColumnarToRowExec(nativeWrite: CometWriteFilesExec) => nativeWrite case c @ ColumnarToRowExec(child) if hasCometNativeChild(child) => val op = createColumnarToRowExec(child) 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..c9cd523f58 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometWriteFiles.scala @@ -0,0 +1,164 @@ +/* + * 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.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 = { + 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 = op.getTagValue(CometExecRule.WRITE_OUTPUT_PATH) 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) + + val outputPath = op.getTagValue(CometExecRule.WRITE_OUTPUT_PATH).get + 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, op.child) + + 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..05de265c71 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala @@ -0,0 +1,298 @@ +/* + * 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.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 com.google.protobuf.CodedOutputStream + +import org.apache.comet.CometExecIterator +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 child + * The Comet native operator producing the batches to write. + */ +case class CometWriteFilesExec(nativeOp: Operator, child: SparkPlan) + extends CometNativeExec + with ShimCometWriteFilesExec { + + override def originalPlan: SparkPlan = child + + 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(serialize(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() + val numPartitions = childRDD.getNumPartitions + val numOutputCols = child.output.length + val childSchema = CometUtils.fromAttributes(child.output) + val nativeOpTemplate = 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. + val dataColumnNames = description.dataColumns.map(_.name) + assert( + dataColumnNames.length == numOutputCols, + s"Expected ${dataColumnNames.length} data columns to write but the child produces " + + s"$numOutputCols") + + childRDD.mapPartitionsInternal { batches => + 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 writeFile = sparkPartitionId == 0 || batches.hasNext + + val writtenFile = if (writeFile) { + 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( + nativeOpTemplate, + filePath, + dataColumnNames, + batches, + childSchema, + numOutputCols, + numPartitions, + sparkPartitionId) + recordRows(statsTrackers, filePath, rowsWritten) + statsTrackers.foreach(_.closeFile(filePath)) + Some(filePath) + } else { + // Drain so the child's native execution completes and releases its resources. + batches.foreach(_.close()) + None + } + + val (taskCommitMessage, taskCommitTime) = Utils.timeTakenMs { + committer.commitTask(taskAttemptContext) + } + logDebug( + s"Task ${taskAttemptContext.getTaskAttemptID} committed " + + writtenFile.getOrElse("no file")) + + 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( + template: Operator, + filePath: String, + dataColumnNames: Seq[String], + batches: Iterator[ColumnarBatch], + childSchema: StructType, + numOutputCols: Int, + numPartitions: Int, + partitionId: Int): Long = { + val parquetWriter = template.getParquetWriter.toBuilder + .setOutputPath(filePath) + .clearColumnNames() + .addAllColumnNames(dataColumnNames.asJava) + .build() + val taskOp = template.toBuilder.setParquetWriter(parquetWriter).build() + + var rowsWritten = 0L + val countingBatches = batches.map { batch => + rowsWritten += batch.numRows() + batch + } + + val execIterator = new CometExecIterator( + CometExec.newIterId, + CometArrowStream.inputObjects(countingBatches, childSchema, nodeName), + numOutputCols, + serialize(taskOp), + CometMetricNode.fromCometPlan(this), + numPartitions, + partitionId, + None, + 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. + */ + private def recordRows( + statsTrackers: Seq[WriteTaskStatsTracker], + filePath: String, + count: Long): Unit = { + statsTrackers.filterNot(_.isInstanceOf[BasicWriteTaskStatsTracker]).foreach { tracker => + 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) { + statsTrackers.foreach(_.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) + } + + private def serialize(op: Operator): Array[Byte] = { + val bytes = new Array[Byte](op.getSerializedSize) + val codedOutput = CodedOutputStream.newInstance(bytes) + op.writeTo(codedOutput) + codedOutput.checkNoSpaceLeft() + bytes + } +} 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..9731424a03 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,40 @@ 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.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 +90,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 +112,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 +146,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 +165,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 +186,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 +211,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 +222,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 +238,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 +471,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 +511,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 +532,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 +567,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 +590,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 +637,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 +665,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 +756,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 +835,33 @@ 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}") } - 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 +871,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 +946,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 +955,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 { From bcfef4b41431767a101705ce7e482146926b6ecf Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 7 Aug 2026 09:48:40 -0600 Subject: [PATCH 2/2] refactor: address review feedback on the native write path Reuse: - Use CometExec.serializeNativePlan and CometExec.getCometIterator instead of hand-rolling protobuf serialization and the CometExecIterator constructor. - Use CometArrowStream.countingIterator instead of a bespoke counting map. Altitude: - originalPlan now points at the WriteFilesExec rather than the child. CometExecRule copies originalPlan's logical link onto every CometExec, so pointing it at the child linked the write node to the child's logical plan; AQE then mistook it for the child's query stage and re-wrapped it in a second WriteFilesExec. That was the actual cause of the nested-write problem, so the collapse guard in CometExecRule is no longer needed and is removed. This also fixes RevertNativeForTransitionHeavyStages, which reverts a CometExec via originalPlan.withNewChildren and would otherwise have dropped the write. - DataWritingCommandExec joins the never-replaced list. It is deliberately left in the plan for a fully native write, so tagging it reported an accelerated write as a fallback in extended explain and skewed the accelerated-operator count. - Delete the ColumnarToRowExec(CometWriteFilesExec) arm in EliminateRedundantTransitions. Spark's ensureOutputsRowBased already special-cases a DataWritingCommandExec over a V1WriteCommand with a columnar child under plannedWrite, so no transition is ever inserted and the arm was dead. - Move the Spark 4.0+ gate from CometExecRule into CometWriteFiles.getSupportLevel, where every other operator's version and shape gating lives. This also gives users on 3.x a reason whether or not they opted in. Efficiency: - The write task closure no longer captures the exec node. Per-task work moves to object CometWriteFilesExec, taking a NativeWriteTask container built on the driver, so the task binary no longer carries nativeOp plus the whole converted child subtree (each node of which holds its own non-transient protobuf). Mirrors how Spark's WriteFilesExec delegates to the static FileFormatWriter.executeTask, and hoists CometMetricNode.fromCometPlan to the driver as CometNativeExec already does. - recordRows loops per tracker on the outside, so the per-row inner loop has one receiver and no per-row closure allocation. Simplification: - createOperatorIncompatConfig takes Option[String]; ConfigBuilder mutates in place, so the Seq destructuring was rebinding the same object. - Flatten the writeFile/writtenFile Option, which only fed a debug message. Also fixes the scalafix CI failure: FileCommitProtocol was imported only for a scaladoc link, which RemoveUnused does not count as a use. Adds a regression assertion that a fully native write leaves no fallback reason on DataWritingCommandExec. --- .../scala/org/apache/comet/CometConf.scala | 12 +- .../apache/comet/rules/CometExecRule.scala | 33 +-- .../rules/EliminateRedundantTransitions.scala | 9 +- .../serde/operator/CometWriteFiles.scala | 30 +- .../spark/sql/comet/CometWriteFilesExec.scala | 275 ++++++++++-------- .../parquet/CometParquetWriterSuite.scala | 15 +- 6 files changed, 200 insertions(+), 174 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6f1b33b40d..3dae1c02e1 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -876,7 +876,7 @@ object CometConf extends ShimCometConf { // already enabled the experimental writer. createOperatorIncompatConfig( "WriteFilesExec", - Seq(getOperatorAllowIncompatConfigKey("DataWritingCommandExec"))) + Some(getOperatorAllowIncompatConfigKey("DataWritingCommandExec"))) /** Create a config to enable a specific operator */ private def createExecEnabledConfig( @@ -903,18 +903,16 @@ object CometConf extends ShimCometConf { private def createOperatorIncompatConfig( name: String, - alternatives: Seq[String] = Seq.empty): ConfigEntry[Boolean] = { + alternative: Option[String] = None): ConfigEntry[Boolean] = { val configKey = getOperatorAllowIncompatConfigKey(name) val envVar = configKeyToEnvVar(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") - val withAlts = alternatives match { - case Nil => builder - case alt +: more => builder.withAlternative(alt, more: _*) - } - withAlts.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 949c9a5fe9..fa6d4cfb65 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -295,29 +295,12 @@ case class CometExecRule(session: SparkSession) case op if shouldApplySparkToColumnar(conf, op) => convertToComet(op, CometSparkToColumnarExec).getOrElse(op) - // AQE re-plans the write command's child, which re-inserts a WriteFilesExec above the node - // Comet already converted on a previous pass. Collapse it instead of wrapping a second - // native write around the first: nested writes would write the data twice, and the inner - // node's empty output would be mistaken for a zero-column schema. - case w: WriteFilesExec if w.child.isInstanceOf[CometWriteFilesExec] => - w.child - // Replace only the per-task write with Comet, leaving DataWritingCommandExec (and therefore - // Spark's commit protocol, stats trackers and SaveMode handling) in place. Gated to Spark - // 4.0+: `V1WritesUtils.getWriteFilesOpt` matches the `WriteFilesExecBase` trait there, which - // is what makes Spark route the write through this node. 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. - case w: WriteFilesExec if isSpark40Plus => + // 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) - // On Spark 3.x the write node cannot be replaced at all, so explain why rather than leaving - // a user who opted into native writes with no fallback reason. Only reported when they asked - // for the feature; otherwise WriteFilesExec falls through to the never-replaced list below. - case w: WriteFilesExec if CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.get(conf) => - withFallbackReason(w, "Native Parquet writes require Spark 4.0 or later") - w - // For AQE broadcast stage on a Comet broadcast exchange case s @ BroadcastQueryStageExec(_, _: CometBroadcastExchangeExec, _) => convertToComet(s, CometExchangeSink).getOrElse(s) @@ -392,12 +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. - // A WriteFilesExec that reaches this point was already offered to CometWriteFiles by - // the case above (or skipped on Spark 3.x, where the node cannot be replaced at all), - // so it has a fallback reason already and must not be tagged again. + // 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 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 ac2b4252cd..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, CometPlan, CometSparkToColumnarExec, CometWriteFilesExec} +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,13 +87,6 @@ case class EliminateRedundantTransitions(session: SparkSession) // and CometSparkToColumnarExec sparkToColumnar.child } - // Remove unnecessary transition for native writes. CometWriteFilesExec reports - // supportsColumnar (it is a CometExec) so that its child stays columnar, but Spark's - // transition rules then wrap it for the row-based DataWritingCommandExec above. That - // wrapper has to go: Spark drives the write through executeWrite on this exact node, and - // ColumnarToRowExec does not implement doExecuteWrite. - case ColumnarToRowExec(nativeWrite: CometWriteFilesExec) => - 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/CometWriteFiles.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometWriteFiles.scala index c9cd523f58..ef1cc335c6 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometWriteFiles.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometWriteFiles.scala @@ -29,7 +29,7 @@ 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.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} @@ -54,6 +54,14 @@ object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { 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")) } @@ -62,7 +70,7 @@ object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { // 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 = op.getTagValue(CometExecRule.WRITE_OUTPUT_PATH) match { + val outputPath = outputPathOf(op) match { case Some(path) => path case None => return Unsupported(Some("Only InsertIntoHadoopFsRelationCommand writes are supported")) @@ -130,11 +138,13 @@ object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { .newBuilder() .setCompression(codec) - val outputPath = op.getTagValue(CometExecRule.WRITE_OUTPUT_PATH).get - val hadoopConf = op.session.sessionState.newHadoopConfWithOptions(op.options) - NativeConfig - .extractObjectStoreOptions(hadoopConf, URI.create(outputPath)) - .foreach { case (key, value) => writerOpBuilder.putObjectStoreOptions(key, value) } + // 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 @@ -146,7 +156,11 @@ object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { } override def createExec(nativeOp: Operator, op: WriteFilesExec): CometNativeExec = - CometWriteFilesExec(nativeOp, op.child) + 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 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 index 05de265c71..6e899c54f8 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala @@ -26,6 +26,7 @@ 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 @@ -39,9 +40,6 @@ import org.apache.spark.sql.types.StructType 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 import org.apache.comet.shims.ShimCometWriteFilesExec @@ -62,22 +60,28 @@ import org.apache.comet.shims.ShimCometWriteFilesExec * 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 + * `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, child: SparkPlan) +case class CometWriteFilesExec( + nativeOp: Operator, + override val originalPlan: SparkPlan, + child: SparkPlan) extends CometNativeExec with ShimCometWriteFilesExec { - override def originalPlan: SparkPlan = child - override def nodeName: String = "CometWriteFiles" override lazy val metrics: Map[String, SQLMetric] = Map( @@ -85,7 +89,8 @@ case class CometWriteFilesExec(nativeOp: Operator, child: SparkPlan) "bytes_written" -> SQLMetrics.createSizeMetric(sparkContext, "written data"), "rows_written" -> SQLMetrics.createMetric(sparkContext, "number of written rows")) - override def serializedPlanOpt: SerializedPlan = SerializedPlan(Some(serialize(nativeOp))) + override def serializedPlanOpt: SerializedPlan = + SerializedPlan(Some(CometExec.serializeNativePlan(nativeOp))) override def withNewChildInternal(newChild: SparkPlan): SparkPlan = copy(child = newChild) @@ -105,88 +110,115 @@ case class CometWriteFilesExec(nativeOp: Operator, child: SparkPlan) val jobTrackerID = SparkHadoopWriterUtils.createJobTrackerID(new Date()) val childRDD = child.executeColumnar() - val numPartitions = childRDD.getNumPartitions - val numOutputCols = child.output.length - val childSchema = CometUtils.fromAttributes(child.output) - val nativeOpTemplate = 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. - val dataColumnNames = description.dataColumns.map(_.name) + + // 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( - dataColumnNames.length == numOutputCols, - s"Expected ${dataColumnNames.length} data columns to write but the child produces " + - s"$numOutputCols") + 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 => - 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 writeFile = sparkPartitionId == 0 || batches.hasNext - - val writtenFile = if (writeFile) { - 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( - nativeOpTemplate, - filePath, - dataColumnNames, - batches, - childSchema, - numOutputCols, - numPartitions, - sparkPartitionId) - recordRows(statsTrackers, filePath, rowsWritten) - statsTrackers.foreach(_.closeFile(filePath)) - Some(filePath) - } else { - // Drain so the child's native execution completes and releases its resources. - batches.foreach(_.close()) - None - } - - val (taskCommitMessage, taskCommitTime) = Utils.timeTakenMs { - committer.commitTask(taskAttemptContext) - } - logDebug( - s"Task ${taskAttemptContext.getTaskAttemptID} committed " + - writtenFile.getOrElse("no file")) - - 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 + 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 } } @@ -197,37 +229,30 @@ case class CometWriteFilesExec(nativeOp: Operator, child: SparkPlan) * because the native writer consumes its whole input before completing. */ private def writeNatively( - template: Operator, + taskWrite: NativeWriteTask, filePath: String, - dataColumnNames: Seq[String], batches: Iterator[ColumnarBatch], - childSchema: StructType, - numOutputCols: Int, - numPartitions: Int, partitionId: Int): Long = { - val parquetWriter = template.getParquetWriter.toBuilder + val parquetWriter = taskWrite.nativeOp.getParquetWriter.toBuilder .setOutputPath(filePath) .clearColumnNames() - .addAllColumnNames(dataColumnNames.asJava) + .addAllColumnNames(taskWrite.dataColumnNames.asJava) .build() - val taskOp = template.toBuilder.setParquetWriter(parquetWriter).build() + val taskOp = taskWrite.nativeOp.toBuilder.setParquetWriter(parquetWriter).build() var rowsWritten = 0L - val countingBatches = batches.map { batch => - rowsWritten += batch.numRows() - batch - } - - val execIterator = new CometExecIterator( - CometExec.newIterId, - CometArrowStream.inputObjects(countingBatches, childSchema, nodeName), - numOutputCols, - serialize(taskOp), - CometMetricNode.fromCometPlan(this), - numPartitions, + 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, - None, - Seq.empty) + broadcastedHadoopConfForEncryption = None, + encryptedFilePaths = Seq.empty) try { // The native writer emits no batches; draining performs the write. @@ -249,21 +274,27 @@ case class CometWriteFilesExec(nativeOp: Operator, child: SparkPlan) * 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.filterNot(_.isInstanceOf[BasicWriteTaskStatsTracker]).foreach { tracker => - 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) { - statsTrackers.foreach(_.newRow(filePath, InternalRow.empty)) - i += 1 + 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 + } } } @@ -287,12 +318,4 @@ case class CometWriteFilesExec(nativeOp: Operator, child: SparkPlan) new TaskAttemptContextImpl(hadoopConf, taskAttemptId) } - - private def serialize(op: Operator): Array[Byte] = { - val bytes = new Array[Byte](op.getSerializedSize) - val codedOutput = CodedOutputStream.newInstance(bytes) - op.writeTo(codedOutput) - codedOutput.checkNoSpaceLeft() - bytes - } } 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 9731424a03..260199652c 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -34,11 +34,12 @@ 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, 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.{CometConf, CometExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, SchemaGenOptions} @@ -848,6 +849,18 @@ class CometParquetWriterSuite extends CometTestBase { assert( nativeWriteCount == 1, 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 assertNoCometWriteFilesExec(plan: SparkPlan): Unit = {