From 27f67828e88685eca326cda798b24b3a1ced894b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 7 Sep 2026 07:51:53 -0600 Subject: [PATCH 1/9] feat: hook native Parquet writes into Spark's WriteFilesExec seam on Spark 4.0+ Native writes replace the whole DataWritingCommandExec, which means InsertIntoHadoopFsRelationCommand.run never runs. Everything that method does has to be re-implemented inside CometNativeWriteExec: a hardcoded SQLHadoopMapReduceCommitProtocol (so spark.sql.sources.commitProtocolClass is ignored), dynamicPartitionOverwrite pinned to false, a hand-ported copy of the SaveMode logic, a bespoke commit-message accumulator, and its own commitJob call. Most of the open native-writer issues are symptoms of that one decision rather than independent defects. On Spark 4.0+, V1WritesUtils.getWriteFilesOpt matches the WriteFilesExecBase trait (introduced in 4.0 precisely for this), so a Comet node that extends it gets driven through FileFormatWriter.executeWrite -> SparkPlan.executeWrite -> doExecuteWrite, and Spark keeps ownership of everything above the per-task write. Spark 3.x has no such trait: getWriteFilesOpt matches the concrete WriteFilesExec case class, a Comet node there would not be found, and Spark would silently take FileFormatWriter's non-planned, row-based branch. So the new seam is additive. CometDataWritingCommand and CometNativeWriteExec are kept unchanged and remain the 3.4/3.5 path; CometExecRule picks the path by version and the two never both fire. The legacy path goes away with Spark 3.x support. Add: - CometWriteFilesExec, overriding doExecuteWrite and mirroring FileFormatWriter.executeTask for the parts Comet must do itself: build the TaskAttemptContext, ask the commit protocol for a path, run the native writer, drive the stats trackers, commit or abort. Plus the CometWriteFiles serde and a two-line ShimCometWriteFilesExec in spark-4.x / spark-3.x. - File paths come from FileCommitProtocol.newTaskTempFile and are used verbatim, so names match Spark's part---c000..parquet and committers that track individual files (S3A magic, streaming manifest) work. - Column names, nullability and field IDs come from WriteJobDescription.dataColumns rather than the query output, so INSERT INTO t SELECT a+1 writes the target column's name. - Byte and row counts come from BasicWriteTaskStatsTracker, which stats files through the FileSystem API and is therefore correct on HDFS. - ParquetWriter proto: work_dir is now optional. When set (3.x) the native writer derives the file name as before; when unset (4.0+) output_path is the exact file to write and is used verbatim. - On 4.0+ the opt-in moves to spark.comet.operator.WriteFilesExec .allowIncompatible, with the old DataWritingCommandExec key kept as a deprecated alternative. isOperatorAllowIncompat now resolves alternatives, which the planner's by-name lookup previously bypassed. AQE re-plans the write command's child and re-inserts a WriteFilesExec above the node Comet already converted; leaving DataWritingCommandExec in place means Comet no longer has to guard against the resulting nested native writes. --- benchmarks/pyspark/run_all_benchmarks.sh | 2 + docs/source/user-guide/latest/installation.md | 4 +- docs/source/user-guide/latest/operators.md | 7 +- .../src/execution/operators/parquet_writer.rs | 89 ++++- native/core/src/execution/planner.rs | 6 +- native/proto/src/proto/operator.proto | 15 +- .../scala/org/apache/comet/CometConf.scala | 40 +- .../apache/comet/rules/CometExecRule.scala | 62 ++- .../serde/operator/CometWriteFiles.scala | 197 +++++++++ .../spark/sql/comet/CometWriteFilesExec.scala | 343 ++++++++++++++++ .../comet/shims/ShimCometWriteFilesExec.scala | 37 ++ .../comet/shims/ShimCometWriteFilesExec.scala | 38 ++ .../parquet/CometParquetWriterSuite.scala | 378 ++++++++++++++++-- .../sql/comet/CometTaskMetricsSuite.scala | 12 +- 14 files changed, 1147 insertions(+), 83 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/serde/operator/CometWriteFiles.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/benchmarks/pyspark/run_all_benchmarks.sh b/benchmarks/pyspark/run_all_benchmarks.sh index 4c23951e0cc..0eafffe640d 100755 --- a/benchmarks/pyspark/run_all_benchmarks.sh +++ b/benchmarks/pyspark/run_all_benchmarks.sh @@ -78,6 +78,7 @@ $SPARK_HOME/bin/spark-submit \ --conf spark.memory.offHeap.enabled=true \ --conf spark.memory.offHeap.size=16g \ --conf spark.comet.enabled=true \ + --conf spark.comet.operator.WriteFilesExec.allowIncompatible=true \ --conf spark.comet.operator.DataWritingCommandExec.allowIncompatible=true \ --conf spark.comet.parquet.write.enabled=true \ --conf spark.comet.explain.fallback.log.enabled=true \ @@ -106,6 +107,7 @@ $SPARK_HOME/bin/spark-submit \ --conf spark.memory.offHeap.enabled=true \ --conf spark.memory.offHeap.size=16g \ --conf spark.comet.enabled=true \ + --conf spark.comet.operator.WriteFilesExec.allowIncompatible=true \ --conf spark.comet.operator.DataWritingCommandExec.allowIncompatible=true \ --conf spark.comet.parquet.write.enabled=true \ --conf spark.comet.explain.fallback.log.enabled=true \ diff --git a/docs/source/user-guide/latest/installation.md b/docs/source/user-guide/latest/installation.md index 7fc4a81d897..06df8340404 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 3a18d86606d..fdd14fb1e34 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -116,9 +116,10 @@ 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` | ⚠️ | Spark 4.0+. Experimental native Parquet writes, disabled by default (opt-in). Non-partitioned, non-bucketed writes only, and not when `spark.sql.files.maxRecordsPerFile` is set. | +| `DataWritingCommandExec` | ⚠️ | Spark 3.4/3.5 only. Experimental native Parquet writes, disabled by default (opt-in). Replaced by `WriteFilesExec` on Spark 4.0+ and removed with Spark 3.x support. | ## Python and UDF diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index 85b76489218..187a84fc5cc 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -220,10 +220,14 @@ impl ParquetWriter { pub struct ParquetWriterExec { /// Input execution plan input: Arc, - /// Output file path (final destination) + /// Where this task writes. When `work_dir` is set (the Spark 3.x `CometNativeWriteExec` + /// path) this is the write's output directory and is unused; the file name is derived from + /// `work_dir`. Otherwise (Spark 4.0+, `CometWriteFilesExec`) it is the exact path of the file + /// to write, chosen by the JVM commit protocol and used verbatim - this operator then never + /// derives file names of its own. output_path: String, - /// Working directory for temporary files (used by FileCommitProtocol) - work_dir: String, + /// Working directory for temporary files (used by FileCommitProtocol). Spark 3.x only. + work_dir: Option, /// Job ID for tracking this write operation job_id: Option, /// Task attempt ID for this specific task @@ -250,7 +254,7 @@ impl ParquetWriterExec { pub fn try_new( input: Arc, output_path: String, - work_dir: String, + work_dir: Option, job_id: Option, task_attempt_id: Option, compression: ParquetCompression, @@ -510,15 +514,18 @@ impl ExecutionPlan for ParquetWriterExec { Arc::new(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) + // Spark 4.0+ hands over the exact file to write, chosen by the JVM commit protocol. + // Spark 3.x hands over a working directory instead and expects the writer to name the + // file; that branch goes away with Spark 3.x support. + let part_file = match &work_dir { + None => self.output_path.clone(), + Some(work_dir) => match task_attempt_id { + Some(attempt_id) => format!( + "{}/part-{:05}-{:05}.parquet", + work_dir, self.partition_id, attempt_id + ), + None => format!("{}/part-{:05}.parquet", work_dir, self.partition_id), + }, }; // Configure writer properties @@ -647,6 +654,56 @@ mod tests { ); } + /// Spark 4.0+ hands over the exact file to write rather than a working directory. The writer + /// must use that path verbatim - Spark's commit protocol owns naming and staging, and + /// committers that track individual files depend on the name it chose. + #[tokio::test] + async fn test_parquet_writer_uses_output_path_verbatim_without_work_dir() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, true)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + + let memory_source = MemorySourceConfig::try_new(&[vec![batch]], Arc::clone(&schema), None)?; + let input = Arc::new(DataSourceExec::new(Arc::new(memory_source))); + let temp_dir = tempfile::tempdir()?; + // A name the writer could never have derived itself, matching Spark's convention. + let file_name = "part-00007-11111111-2222-3333-4444-555555555555-c000.parquet"; + let output_path = format!("file://{}/{}", temp_dir.path().display(), file_name); + + let writer = ParquetWriterExec::try_new( + input, + output_path, + None, // work_dir: Spark 4.0+ path + None, + None, + ParquetCompression::None, + // A non-zero partition id must not leak into the file name. + 3, + vec!["id".to_string()], + None, + HashMap::new(), + )?; + + let mut stream = writer.execute(0, SessionContext::new().task_ctx())?; + while stream.try_next().await?.is_some() {} + + let written = temp_dir.path().join(file_name); + assert!( + written.exists(), + "expected the writer to use the given path verbatim, found: {:?}", + std::fs::read_dir(temp_dir.path())? + .filter_map(|e| e.ok().map(|e| e.file_name())) + .collect::>() + ); + + let reader = SerializedFileReader::new(File::open(written)?)?; + assert_eq!(reader.metadata().file_metadata().num_rows(), 3); + + Ok(()) + } + #[tokio::test] async fn test_parquet_writer_preserves_catalyst_schema_in_footer() -> Result<()> { let values = ListArray::from_iter_primitive::([ @@ -689,7 +746,7 @@ mod tests { let writer = ParquetWriterExec::try_new( input, work_dir.clone(), - work_dir, + Some(work_dir), None, None, ParquetCompression::None, @@ -758,7 +815,7 @@ mod tests { let writer = ParquetWriterExec::try_new( input, work_dir.clone(), - work_dir, + Some(work_dir), None, None, ParquetCompression::None, @@ -1018,7 +1075,7 @@ mod tests { let parquet_writer = ParquetWriterExec::try_new( memory_exec, output_path, - work_dir, + Some(work_dir), None, // job_id Some(123), // task_attempt_id ParquetCompression::None, diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 5cfe74aa412..dbf86aeddc1 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1991,11 +1991,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.work_dir.clone(), writer.job_id.clone(), writer.task_attempt_id, codec, diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 7b34f84012e..2f437f5398e 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -842,11 +842,22 @@ message ShuffleWriter { } message ParquetWriter { + // Where this task writes. Two shapes, selected by whether `work_dir` is set: + // + // - Spark 4.0+ (CometWriteFilesExec): `work_dir` is unset and `output_path` is the + // fully-qualified path of the Parquet file to write, set per task 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. + // - Spark 3.x (CometNativeWriteExec): `work_dir` is set and the native writer derives the file + // name from it, the partition id and the task attempt id. `output_path` is the write's + // output directory and is unused natively. Goes away with Spark 3.x support. 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 + // Working directory for temporary files (used by FileCommitProtocol). Spark 3.x only; see + // output_path above. optional string work_dir = 5; // Job ID for tracking this write operation optional string job_id = 6; diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index c015d2d4e66..b6bebfe9b09 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -1005,9 +1005,29 @@ object CometConf extends ShimCometConf { .booleanConf .createWithEnvVarOrDefault("ENABLE_COMET_STRICT_TESTING", false) + /** + * Deprecated alternatives for `spark.comet.operator..allowIncompatible`, keyed by + * operator name. `CometExecRule` resolves these configs by operator name rather than through + * the registered `ConfigEntry`, so `isOperatorAllowIncompat` has to consult the alternatives + * itself - otherwise an old key would read `true` from the entry while the planner saw `false`. + */ + private val operatorIncompatAlternatives = + scala.collection.mutable.Map.empty[String, Seq[String]] + + /** Spark 3.x only: native writes replace the whole `DataWritingCommandExec` there. */ val COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT: ConfigEntry[Boolean] = createOperatorIncompatConfig("DataWritingCommandExec") + /** + * Spark 4.0+ only: native writes replace just the `WriteFilesExec` child, so the opt-in moved + * with the operator. The old `DataWritingCommandExec` key keeps working for anyone who had + * already enabled the experimental writer. + */ + val COMET_OPERATOR_WRITE_FILES_ALLOW_INCOMPAT: ConfigEntry[Boolean] = + createOperatorIncompatConfig( + "WriteFilesExec", + Seq(getOperatorAllowIncompatConfigKey("DataWritingCommandExec"))) + /** Create a config to enable a specific operator */ private def createExecEnabledConfig( exec: String, @@ -1031,15 +1051,21 @@ 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] = Nil): ConfigEntry[Boolean] = { val configKey = getOperatorAllowIncompatConfigKey(name) val envVar = configKeyToEnvVar(configKey) - conf(configKey) + // ConfigBuilder mutates in place, so the result of withAlternative is `builder` itself. + 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) + if (alternatives.nonEmpty) { + operatorIncompatAlternatives.put(name, alternatives) + builder.withAlternative(alternatives.head, alternatives.tail: _*) + } + builder.booleanConf.createWithEnvVarOrDefault(envVar, false) } def isExprEnabled(name: String, conf: SQLConf = SQLConf.get): Boolean = { @@ -1063,7 +1089,11 @@ object CometConf extends ShimCometConf { } def isOperatorAllowIncompat(name: String, conf: SQLConf = SQLConf.get): Boolean = { - getBooleanConf(getOperatorAllowIncompatConfigKey(name), defaultValue = false, conf) + val value = CometConfDeprecations.readWithAlternatives( + conf, + getOperatorAllowIncompatConfigKey(name), + operatorIncompatAlternatives.getOrElse(name, Nil)) + value != null && value.toLowerCase(Locale.ROOT) == "true" } def getOperatorAllowIncompatConfigKey(name: String): String = { 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 db5373e1093..00246ccce30 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -37,7 +37,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffl import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec 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 @@ -106,6 +106,15 @@ 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 [[org.apache.comet.serde.operator.CometWriteFiles]] learns the target filesystem - and, + * by its absence, that a write comes from some other V1 write command. Only used on Spark 4.0+; + * see `CometWriteFilesExec`. + */ + 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`. @@ -389,15 +398,27 @@ case class CometExecRule(session: SparkSession) case op if shouldApplySparkToColumnar(conf, op) => convertToComet(op, CometSparkToColumnarExec).getOrElse(op) + // Spark 4.0+: replace only the per-task write, leaving DataWritingCommandExec - and + // therefore Spark's commit protocol, stats trackers and SaveMode handling - in place. + // `V1WritesUtils.getWriteFilesOpt` matches the `WriteFilesExecBase` trait there, which is + // what lets a Comet node stand in for the write node. See CometWriteFilesExec. + case w: WriteFilesExec if isSpark40Plus => + convertToComet(w, CometWriteFiles).getOrElse(w) + + // Spark 3.x: `getWriteFilesOpt` matches the concrete `WriteFilesExec` case class, so a + // Comet node can never stand in for the write node. Native writes instead replace the whole + // DataWritingCommandExec and re-implement the write framework inside CometNativeWriteExec. + // This path is retained only for 3.4/3.5 and goes away with them. + // // 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] => + if !isSpark40Plus && w.child.isInstanceOf[CometNativeWriteExec] => w.child - case op: DataWritingCommandExec => + case op: DataWritingCommandExec if !isSpark40Plus => convertToComet(op, CometDataWritingCommand).getOrElse(op) // AQE re-fires the Iceberg write planning on every stage materialisation, so a @@ -484,14 +505,21 @@ case class CometExecRule(session: SparkSession) op match { case _: CometPlan | _: AQEShuffleReadExec | _: BroadcastExchangeExec | _: BroadcastQueryStageExec | _: AdaptiveSparkPlanExec | _: ExecutedCommandExec | - _: V2CommandExec | _: WriteFilesExec => + _: V2CommandExec => // 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. + op + case _: WriteFilesExec if !isSpark40Plus => + // On Spark 3.x the write is converted at the DataWritingCommandExec above, which + // unwraps WriteFilesExec inside convertToComet. Tagging it 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. On + // 4.0+ the node was offered to CometWriteFiles above and already carries a reason. + op + case _: DataWritingCommandExec if isSpark40Plus => + // On Spark 4.0+ DataWritingCommandExec is deliberately left in the plan even for a + // fully native write - Comet replaces only its WriteFilesExec child - so tagging it + // would report an accelerated write as a fallback. op case _ => // The operator was not converted to a Comet plan and no serde handler claimed it, so @@ -513,6 +541,19 @@ 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. Only the Spark 4.0+ path consults this tag. + if (isSpark40Plus) { + 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 @@ -802,7 +843,8 @@ case class CometExecRule(session: SparkSession) // its child (e.g., CometNativeScanExec, or a CometProject over an AQEShuffleRead) // needs its own serialization. Reset the flag so children can start their own native // execution blocks. - if (op.isInstanceOf[CometNativeWriteExec] || op.isInstanceOf[CometIcebergWriteExec]) { + if (op.isInstanceOf[CometNativeWriteExec] || op.isInstanceOf[CometIcebergWriteExec] || + op.isInstanceOf[CometWriteFilesExec]) { firstNativeOp = true } 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 00000000000..ca81f15b688 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometWriteFiles.scala @@ -0,0 +1,197 @@ +/* + * 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.catalyst.util.CaseInsensitiveMap +import org.apache.spark.sql.comet.{CometNativeExec, CometWriteFilesExec} +import org.apache.spark.sql.execution.datasources.WriteFilesExec +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.{CometConf, ConfigEntry} +import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, withFallbackReason} +import org.apache.comet.objectstore.NativeConfig +import org.apache.comet.rules.CometExecRule +import org.apache.comet.serde.{CometOperatorSerde, Incompatible, OperatorOuterClass, SupportLevel, Unsupported} +import org.apache.comet.serde.OperatorOuterClass.Operator + +/** + * Serde for Spark's `WriteFilesExec`, replacing the per-task Parquet write with Comet's native + * writer while leaving the surrounding write framework (commit protocol, stats trackers, SaveMode + * handling, `_SUCCESS`) to Spark. See [[CometWriteFilesExec]] for how the two fit together. + */ +object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { + + private val supportedCompressionCodecs = + Set("none", "uncompressed", "snappy", "lz4", "zstd", "gzip") + + override def enabledConfig: Option[ConfigEntry[Boolean]] = + Some(CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED) + + // Native writes require Arrow-formatted input data. If the query falls back to Spark + // (e.g., due to unsupported complex types), the write must also fall back. + override def requiresNativeChildren: Boolean = true + + override def getSupportLevel(op: WriteFilesExec): SupportLevel = { + // `V1WritesUtils.getWriteFilesOpt` matches the `WriteFilesExecBase` trait on Spark 4.0+, which + // is what makes Spark route the write through CometWriteFilesExec. Spark 3.x matches the + // concrete `WriteFilesExec` case class instead, so a Comet node would be silently ignored and + // the write would fall into FileFormatWriter's non-planned, row-based branch. Native writes + // there go through CometDataWritingCommand instead; `CometExecRule` never offers a + // WriteFilesExec to this serde on 3.x, so this guard is only a safety net. + if (!isSpark40Plus) { + return Unsupported(Some("Native Parquet writes require Spark 4.0 or later")) + } + + if (!op.fileFormat.isInstanceOf[ParquetFileFormat]) { + return Unsupported(Some("Only Parquet writes are supported")) + } + + // The write node does not carry the output path, so CometExecRule tags it from the enclosing + // InsertIntoHadoopFsRelationCommand. An absent tag means this write belongs to some other V1 + // write command (a Hive insert, for example) whose semantics Comet has not been verified + // against, so decline it. + val outputPath = outputPathOf(op) match { + case Some(path) => path + case None => + return Unsupported(Some("Only InsertIntoHadoopFsRelationCommand writes are supported")) + } + + if (!outputPath.startsWith("file:") && !outputPath.startsWith("hdfs:")) { + return Unsupported(Some("Supported output filesystems: local, HDFS")) + } + + if (op.bucketSpec.isDefined) { + return Unsupported(Some("Bucketed writes are not supported")) + } + + if (op.partitionColumns.nonEmpty || op.staticPartitions.nonEmpty) { + // This also declines dynamic partition overwrite. `InsertIntoHadoopFsRelationCommand` only + // sets `dynamicPartitionOverwrite` when `staticPartitions.size < partitionColumns.length`, + // which implies partition columns, so a dynamic overwrite always lands here. + return Unsupported(Some("Partitioned writes are not supported")) + } + + if (rollsFilesByRecordCount(op)) { + return Unsupported( + Some("Writes with spark.sql.files.maxRecordsPerFile set 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 scanOperator = NativeWriteUtils.buildFfiScan(op.child, op.id) match { + case Some(scan) => scan + case None => + withFallbackReason(op, "Cannot serialize data types for native write") + return None + } + + 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`, `column_names` and `output_schema` are filled in per task by + // CometWriteFilesExec: the path comes from the commit protocol and the columns from + // WriteJobDescription.dataColumns, neither of which is known at planning time. + val writerOpBuilder = OperatorOuterClass.ParquetWriter + .newBuilder() + .setCompression(codec) + + // getSupportLevel already declined the write if the tag is absent, so this cannot be empty. + outputPathOf(op).foreach { outputPath => + val hadoopConf = op.session.sessionState.newHadoopConfWithOptions(op.options) + NativeConfig + .extractObjectStoreOptions(hadoopConf, URI.create(outputPath)) + .foreach { case (key, value) => writerOpBuilder.putObjectStoreOptions(key, value) } + } + + Some( + Operator + .newBuilder() + .setPlanId(op.id) + .addChildren(scanOperator) + .setParquetWriter(writerOpBuilder.build()) + .build()) + } + + override def createExec(nativeOp: Operator, op: WriteFilesExec): CometNativeExec = + CometWriteFilesExec(nativeOp, originalPlan = op, child = op.child) + + /** The write's output path, recorded on the node by `CometExecRule`. */ + private def outputPathOf(op: WriteFilesExec): Option[String] = + op.getTagValue(CometExecRule.WRITE_OUTPUT_PATH) + + /** + * Whether Spark would roll to a new file every N rows within a task. + * + * `SingleDirectoryDataWriter.write` rolls a new file once `maxRecordsPerFile` rows have been + * written, incrementing the `-c$fileCounter%03d` suffix. Comet's writer asks the commit + * protocol for one file per task and always writes `-c000`, so a task that should have produced + * several files would produce one oversized file instead - a different layout than Spark's own + * writer with no error. Decline the write rather than silently ignore the setting. + */ + private def rollsFilesByRecordCount(op: WriteFilesExec): Boolean = { + // Same precedence as FileFormatWriter.write: the `maxRecordsPerFile` write option wins over + // spark.sql.files.maxRecordsPerFile. Options are matched case-insensitively there. + val maxRecordsPerFile = CaseInsensitiveMap(op.options) + .get("maxRecordsPerFile") + .map(_.toLong) + .getOrElse(SQLConf.get.maxRecordsPerFile) + maxRecordsPerFile > 0 + } + + 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/CometWriteFilesExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala new file mode 100644 index 00000000000..1beeb0c8537 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala @@ -0,0 +1,343 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet + +import java.util.Date + +import scala.jdk.CollectionConverters._ + +import org.apache.hadoop.mapreduce.{TaskAttemptContext, TaskAttemptID, TaskID, TaskType} +import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl +import org.apache.spark.TaskContext +import org.apache.spark.internal.Logging +import org.apache.spark.internal.io.{FileCommitProtocol, FileNameSpec, SparkHadoopWriterUtils} +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.comet.execution.arrow.CometArrowStream +import org.apache.spark.sql.comet.util.{Utils => CometUtils} +import org.apache.spark.sql.connector.write.WriterCommitMessage +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.{BasicWriteTaskStatsTracker, ExecutedWriteSummary, WriteFilesSpec, WriteJobDescription, WriteTaskResult, WriteTaskStatsTracker} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.Utils + +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.schema2Proto +import org.apache.comet.shims.ShimCometWriteFilesExec + +/** + * Comet's replacement for Spark's `WriteFilesExec`: writes Parquet files natively for one task, + * and nothing else. + * + * Everything around the per-task write stays on Spark's side. Because this node extends + * `WriteFilesExecBase` (see [[ShimCometWriteFilesExec]]), `V1WritesUtils.getWriteFilesOpt` finds + * it, so `InsertIntoHadoopFsRelationCommand.run` and `FileFormatWriter` continue to own: + * + * - SaveMode semantics and the delete-before-overwrite of the target directory + * - committer instantiation (including `spark.sql.sources.commitProtocolClass`), `setupJob`, + * `commitJob`/`abortJob`, `onTaskCommit`, and the `_SUCCESS` marker + * - dynamic partition overwrite and custom partition locations + * - `WriteJobStatsTracker` aggregation, SQL metrics, and catalog statistics/cache refresh + * + * This node mirrors `FileFormatWriter.executeTask` for the parts Comet must do itself: build the + * `TaskAttemptContext`, ask the commit protocol where to write, run the native writer, drive the + * stats trackers, and commit or abort the task. Notably the path handed back by + * `FileCommitProtocol.newTaskTempFile` is used verbatim, so staging directories, task-attempt + * isolation under speculation, and committers that track individual files all behave as they do + * for Spark's own writer. + * + * @param nativeOp + * Template for the native write plan. `output_path` is a placeholder here and is replaced per + * task with the path the commit protocol chose. + * @param originalPlan + * The `WriteFilesExec` this node replaced. This must be the write node rather than the data + * subtree: `CometExecRule` copies `originalPlan`'s logical link onto every `CometExec`, and + * pointing it at the child would link the write node to the child's logical plan, which makes + * AQE mistake it for the child's query stage and re-wrap it in a second `WriteFilesExec`. + * @param child + * The Comet native operator producing the batches to write. + */ +case class CometWriteFilesExec( + nativeOp: Operator, + override val originalPlan: SparkPlan, + child: SparkPlan) + extends CometNativeExec + with ShimCometWriteFilesExec { + + override def nodeName: String = "CometWriteFiles" + + override lazy val metrics: Map[String, SQLMetric] = Map( + "files_written" -> SQLMetrics.createMetric(sparkContext, "number of written data files"), + "bytes_written" -> SQLMetrics.createSizeMetric(sparkContext, "written data"), + "rows_written" -> SQLMetrics.createMetric(sparkContext, "number of written rows")) + + override def serializedPlanOpt: SerializedPlan = + SerializedPlan(Some(CometExec.serializeNativePlan(nativeOp))) + + override def withNewChildInternal(newChild: SparkPlan): SparkPlan = copy(child = newChild) + + /** + * Spark drives this node through `executeWrite`, never `execute`. `WriteFilesExecBase` already + * throws for `doExecute`, but `CometExec` widens it to a public member that returns a + * `ColumnarToRowExec` result, so the conflict has to be resolved explicitly here. + */ + override def doExecute(): RDD[InternalRow] = + throw new UnsupportedOperationException(s"$nodeName does not support doExecute") + + override protected def doExecuteWrite( + writeFilesSpec: WriteFilesSpec): RDD[WriterCommitMessage] = { + val description = writeFilesSpec.description + val committer = writeFilesSpec.committer + // Same identifier scheme as FileFormatWriter, so committers that parse the job ID agree. + val jobTrackerID = SparkHadoopWriterUtils.createJobTrackerID(new Date()) + + val childRDD = child.executeColumnar() + + // SPARK-23271 (defensive): a zero-partition input would spawn no task and therefore write no + // file at all, so the output directory would carry no schema for readers. Spark's own + // WriteFilesExec swaps in a dummy single-partition RDD for exactly this case. In practice + // CometWriteFiles.requiresNativeChildren rules out the sources (LocalTableScan) that produce + // a zero-partition RDD, but the swap is kept to match Spark's semantics if that ever changes. + val writeRDD = if (childRDD.getNumPartitions == 0) { + sparkContext.parallelize(Seq.empty[ColumnarBatch], 1) + } else { + childRDD + } + + // 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. + // The write's target schema, not the query output's: 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, nullability and Parquet field IDs. + val dataSchema = CometUtils.fromAttributes(description.dataColumns) + + val taskWrite = NativeWriteTask( + nativeOp = nativeOp, + dataColumnNames = dataSchema.fields.map(_.name).toSeq, + outputSchema = schema2Proto( + dataSchema.fields.toIndexedSeq, + Some(conf.getConf(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED))), + childSchema = CometUtils.fromAttributes(child.output), + numPartitions = writeRDD.getNumPartitions, + nativeMetrics = CometMetricNode.fromCometPlan(this), + nodeName = nodeName) + + assert( + taskWrite.dataColumnNames.length == child.output.length, + s"Expected ${taskWrite.dataColumnNames.length} data columns to write but the child " + + s"produces ${child.output.length}") + + writeRDD.mapPartitionsInternal { batches => + CometWriteFilesExec.executeTask(description, committer, jobTrackerID, taskWrite, batches) + } + } +} + +/** + * The per-task state that [[CometWriteFilesExec.executeTask]] needs, resolved on the driver. + * + * A plain container rather than a closure over the exec node: it copies only these fields, so the + * enclosing plan tree is not kept alive for the task's lifetime or shipped in the task binary. + */ +private[comet] case class NativeWriteTask( + nativeOp: Operator, + dataColumnNames: Seq[String], + outputSchema: Seq[OperatorOuterClass.SparkStructField], + childSchema: StructType, + numPartitions: Int, + nativeMetrics: CometMetricNode, + nodeName: String) + +object CometWriteFilesExec extends Logging { + + /** + * Write one task's batches natively and commit or abort it, mirroring the structure of + * `FileFormatWriter.executeTask`. + */ + private[comet] def executeTask( + description: WriteJobDescription, + committer: FileCommitProtocol, + jobTrackerID: String, + taskWrite: NativeWriteTask, + batches: Iterator[ColumnarBatch]): Iterator[WriterCommitMessage] = { + val taskCtx = TaskContext.get() + val sparkPartitionId = taskCtx.partitionId() + val taskAttemptContext = createTaskAttemptContext( + description, + jobTrackerID, + taskCtx.stageId(), + sparkPartitionId, + // Truncation to Int matches FileFormatWriter: the masked low bits are what the Hadoop + // TaskAttemptID accepts, and uniqueness within a job is preserved by the task ID. + taskCtx.taskAttemptId().toInt & Integer.MAX_VALUE) + + committer.setupTask(taskAttemptContext) + val statsTrackers = description.statsTrackers.map(_.newTaskInstance()) + + try { + // Mirrors FileFormatWriter's EmptyDirectoryDataWriter case: an empty input still writes one + // file from partition 0 so that the output carries the schema, but every other empty + // partition produces no file at all. + val writtenFile = if (sparkPartitionId == 0 || batches.hasNext) { + val ext = description.outputWriterFactory.getFileExtension(taskAttemptContext) + // FileNameSpec's "-c000" suffix reproduces Spark's part---c000..parquet + // naming. The file counter is always 0 until file rolling is supported. + val filePath = + committer.newTaskTempFile(taskAttemptContext, None, FileNameSpec("", "-c000" + ext)) + + statsTrackers.foreach(_.newFile(filePath)) + val rowsWritten = writeNatively(taskWrite, filePath, batches, sparkPartitionId) + recordRows(statsTrackers, filePath, rowsWritten) + statsTrackers.foreach(_.closeFile(filePath)) + filePath + } else { + // Drain so the child's native execution completes and releases its resources. + batches.foreach(_.close()) + "no file" + } + + val (taskCommitMessage, taskCommitTime) = Utils.timeTakenMs { + committer.commitTask(taskAttemptContext) + } + logDebug(s"Task ${taskAttemptContext.getTaskAttemptID} committed $writtenFile") + + Iterator( + WriteTaskResult( + taskCommitMessage, + ExecutedWriteSummary( + // Only non-partitioned writes are supported so far, so no partition paths were + // added. Populating this is part of adding partitioned write support. + updatedPartitions = Set.empty, + stats = statsTrackers.map(_.getFinalStats(taskCommitTime))))) + } catch { + case t: Throwable => + Utils.tryLogNonFatalError(committer.abortTask(taskAttemptContext)) + logError(s"Task ${taskAttemptContext.getTaskAttemptID} aborted: ${t.getMessage}", t) + throw t + } + } + + /** + * Run the native write plan for one task, returning the number of rows written. + * + * The row count is taken on the JVM side as batches are pulled into native code. That is exact + * because the native writer consumes its whole input before completing. + */ + private def writeNatively( + taskWrite: NativeWriteTask, + filePath: String, + batches: Iterator[ColumnarBatch], + partitionId: Int): Long = { + val parquetWriter = taskWrite.nativeOp.getParquetWriter.toBuilder + .setOutputPath(filePath) + .clearColumnNames() + .addAllColumnNames(taskWrite.dataColumnNames.asJava) + .clearOutputSchema() + .addAllOutputSchema(taskWrite.outputSchema.asJava) + .build() + val taskOp = taskWrite.nativeOp.toBuilder.setParquetWriter(parquetWriter).build() + + var rowsWritten = 0L + val countingBatches = + CometArrowStream.countingIterator[ColumnarBatch](batches, b => rowsWritten += b.numRows()) + + val execIterator = CometExec.getCometIterator( + CometArrowStream.inputObjects(countingBatches, taskWrite.childSchema, taskWrite.nodeName), + taskWrite.dataColumnNames.length, + taskOp, + taskWrite.nativeMetrics, + taskWrite.numPartitions, + partitionId, + broadcastedHadoopConfForEncryption = None, + encryptedFilePaths = Seq.empty) + + try { + // The native writer emits no batches; draining performs the write. + while (execIterator.hasNext) { + execIterator.next().close() + } + } finally { + execIterator.close() + } + + rowsWritten + } + + /** + * Report `count` rows to each stats tracker. + * + * `WriteTaskStatsTracker.newRow` is a per-row callback, but the only implementation Spark + * ships, [[BasicWriteTaskStatsTracker]], ignores the row and just counts. Comet has columnar + * batches rather than `InternalRow`s here, so it passes an empty row instead of materializing + * every row just to hand it straight back. A tracker that actually inspects row contents would + * therefore see empty rows, so warn rather than silently report wrong statistics. + * + * The loop is per-tracker on the outside so the hot inner loop has a single receiver and no + * per-row closure; the trackers are independent per-file counters, so their relative + * interleaving carries no meaning. + */ + private def recordRows( + statsTrackers: Seq[WriteTaskStatsTracker], + filePath: String, + count: Long): Unit = { + statsTrackers.foreach { tracker => + if (!tracker.isInstanceOf[BasicWriteTaskStatsTracker]) { + logWarning( + s"${tracker.getClass.getName} receives row counts but not row contents from Comet's " + + "native Parquet writer. Set spark.comet.parquet.write.enabled=false if this tracker " + + "needs to inspect written rows.") + } + var i = 0L + while (i < count) { + tracker.newRow(filePath, InternalRow.empty) + i += 1 + } + } + } + + /** Build the `TaskAttemptContext` exactly as `FileFormatWriter.executeTask` does. */ + private def createTaskAttemptContext( + description: WriteJobDescription, + jobTrackerID: String, + sparkStageId: Int, + sparkPartitionId: Int, + sparkAttemptNumber: Int): TaskAttemptContext = { + val jobId = SparkHadoopWriterUtils.createJobID(jobTrackerID, sparkStageId) + val taskId = new TaskID(jobId, TaskType.MAP, sparkPartitionId) + val taskAttemptId = new TaskAttemptID(taskId, sparkAttemptNumber) + + val hadoopConf = description.serializableHadoopConf.value + hadoopConf.set("mapreduce.job.id", jobId.toString) + hadoopConf.set("mapreduce.task.id", taskAttemptId.getTaskID.toString) + hadoopConf.set("mapreduce.task.attempt.id", taskAttemptId.toString) + hadoopConf.setBoolean("mapreduce.task.ismap", true) + hadoopConf.setInt("mapreduce.task.partition", 0) + + new TaskAttemptContextImpl(hadoopConf, taskAttemptId) + } +} diff --git a/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala new file mode 100644 index 00000000000..f3bb6bbd2a4 --- /dev/null +++ b/spark/src/main/spark-3.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.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 on 3.x go through `CometDataWritingCommand`, which replaces the + * whole `DataWritingCommandExec` instead; `CometExecRule` never offers a `WriteFilesExec` to + * `CometWriteFiles` on 3.x. This shim exists only so that the shared sources compile against 3.x + * and 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 00000000000..63ae713950b --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala @@ -0,0 +1,38 @@ +/* + * 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 this seam is gated to Spark 4.0+ in `CometExecRule` and native writes there go + * through `CometDataWritingCommand` instead. 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 eef77d88246..97b3f417c29 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -19,26 +19,29 @@ package org.apache.comet.parquet -import java.io.File +import java.io.{File, IOException} import scala.jdk.CollectionConverters._ import scala.util.{Random, Using} import org.apache.hadoop.fs.{FileSystem, Path} +import org.apache.hadoop.mapreduce.TaskAttemptContext import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.metadata.CompressionCodecName import org.apache.parquet.hadoop.util.HadoopInputFile import org.apache.parquet.schema.{MessageType, Type} +import org.apache.spark.internal.io.FileCommitProtocol 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, CometNativeWriteExec, 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.SQLHadoopMapReduceCommitProtocol import org.apache.spark.sql.functions.{array, map, struct, when} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, LongType, MapType, Metadata, MetadataBuilder, StringType, StructField, StructType} -import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus +import org.apache.comet.{CometConf, CometExplainInfo} +import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, SchemaGenOptions} class CometParquetWriterSuite extends CometTestBase { @@ -77,7 +80,7 @@ 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", + nativeWriteAllowIncompatKey -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true") { writeWithCometNativeWriteExec(inputPath, outputPath) @@ -99,7 +102,7 @@ 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", + nativeWriteAllowIncompatKey -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true") { val capturedPlan = writeWithCometNativeWriteExec(inputPath, outputPath) @@ -133,8 +136,7 @@ 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", + nativeWriteAllowIncompatKey -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true") { writeWithCometNativeWriteExec(inputPath, outputPath, Some(10)) @@ -363,7 +365,7 @@ class CometParquetWriterSuite extends CometTestBase { withSQLConf( CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", - CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true", + nativeWriteAllowIncompatKey -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", SQLConf.PARQUET_COMPRESSION.key -> codec) { @@ -384,7 +386,7 @@ class CometParquetWriterSuite extends CometTestBase { withSQLConf( CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", - CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true", + nativeWriteAllowIncompatKey -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", SQLConf.PARQUET_COMPRESSION.key -> "snappy") { @@ -409,7 +411,7 @@ class CometParquetWriterSuite extends CometTestBase { withSQLConf( CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", - CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true", + nativeWriteAllowIncompatKey -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", SQLConf.PARQUET_COMPRESSION.key -> "zstd") { @@ -436,7 +438,7 @@ class CometParquetWriterSuite extends CometTestBase { withSQLConf( CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", - CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true", + nativeWriteAllowIncompatKey -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", SQLConf.PARQUET_COMPRESSION.key -> "lz4_raw") { @@ -863,6 +865,255 @@ class CometParquetWriterSuite extends CometTestBase { } } + // --------------------------------------------------------------------------------------------- + // Spark 4.0+ only. These cover behavior that comes from leaving Spark's write framework in + // place, which is only possible where `V1WritesUtils.getWriteFilesOpt` matches the + // `WriteFilesExecBase` trait. See CometWriteFilesExec. + // --------------------------------------------------------------------------------------------- + + test("write creates a _SUCCESS marker") { + assume(isSpark40Plus, "Requires the WriteFilesExec seam") + // 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) + assertHasCometNativeWriteExec(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") { + assume(isSpark40Plus, "Requires the WriteFilesExec seam") + // 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) + assertHasCometNativeWriteExec(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") { + assume(isSpark40Plus, "Requires the WriteFilesExec seam") + // 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) + } + } + + test("dynamic partition overwrite falls back to Spark") { + assume(isSpark40Plus, "Requires the WriteFilesExec seam") + // A dynamic overwrite is a partitioned write, which CometWriteFiles declines - but the + // consequence of getting it wrong is silent data loss across untouched partitions, so assert + // the fallback and the semantics explicitly rather than relying on the partitioning check. + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val original = Seq((1, "a"), (2, "b")).toDF("id", "part") + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + original.write.partitionBy("part").parquet(outputPath) + } + + withNativeWriter { + withSQLConf(SQLConf.PARTITION_OVERWRITE_MODE.key -> "DYNAMIC") { + val replacement = Seq((3, "b")).toDF("id", "part") + val plan = captureWritePlan( + p => replacement.write.mode(SaveMode.Overwrite).partitionBy("part").parquet(p), + outputPath) + assertNoCometNativeWriteExec(plan) + } + } + + // part=a is untouched, part=b is replaced: the defining property of a dynamic overwrite. + checkAnswer(spark.read.parquet(outputPath), Row(1, "a") :: Row(3, "b") :: Nil) + } + } + + test("write with maxRecordsPerFile falls back to Spark") { + assume(isSpark40Plus, "Requires the WriteFilesExec seam") + // Spark's SingleDirectoryDataWriter rolls a new file every maxRecordsPerFile rows, bumping the + // -c suffix. Comet asks the commit protocol for one file per task, so it must decline + // rather than quietly produce a different file layout. + Seq( + "spark.sql.files.maxRecordsPerFile" -> ((df: DataFrame, p: String) => df.write.parquet(p)), + // The write option takes precedence over the conf in FileFormatWriter, so it must be + // honored here too - with the conf left at its default of 0. + "maxRecordsPerFile-option" -> ((df: DataFrame, p: String) => + df.write.option("maxRecordsPerFile", "10").parquet(p))).foreach { case (label, write) => + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val sourcePath = new File(dir, "source.parquet").getAbsolutePath + withNativeWriter { + val df = materializeAsCometSource( + (1 to 100).map(i => (i, s"str_$i")).toDF("id", "name").repartition(1), + sourcePath) + val confs = + if (label == "spark.sql.files.maxRecordsPerFile") Seq(label -> "10") else Seq.empty + withSQLConf(confs: _*) { + val plan = captureWritePlan(p => write(df, p), outputPath) + assertNoCometNativeWriteExec(plan) + } + checkAnswer(spark.read.parquet(outputPath), df) + } + // Spark's writer rolled the 100 rows of the single partition into 10 files of 10 rows. + assert( + listPartFileNames(outputPath).size == 10, + s"$label: expected 10 rolled part files, got ${listPartFileNames(outputPath)}") + } + } + } + + test("empty input still writes a schema-only file (SPARK-23271)") { + assume(isSpark40Plus, "Requires the WriteFilesExec seam") + // An empty input must still leave a schema behind for downstream readers: `spark.read.parquet` + // of the output must see the write's schema, not fail. Comet reaches this in two ways - if the + // native child has one partition producing no batches, the partition-0 branch of executeTask + // writes a metadata-only file; if it produces zero partitions, doExecuteWrite swaps in a dummy + // single-partition RDD to get to the same branch. This test exercises the reachable path + // (filtered Comet scan yielding an empty batch iterator); the zero-partition swap is defensive + // because CometWriteFiles.requiresNativeChildren rules out the sources (LocalTableScan) that + // would otherwise produce a zero-partition RDD. + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val sourcePath = new File(dir, "source.parquet").getAbsolutePath + withNativeWriter { + val empty = materializeAsCometSource( + (1 to 100).map(i => (i, s"str_$i")).toDF("id", "name"), + sourcePath).where("id < 0") + val plan = captureWritePlan(p => empty.write.parquet(p), outputPath) + assertHasCometNativeWriteExec(plan) + + val partFiles = listPartFileNames(outputPath) + assert(partFiles.size == 1, s"Expected exactly one schema-only part file, got $partFiles") + // Reading without an explicit schema is the point: the file must carry it. + val readBack = spark.read.parquet(outputPath) + assert(readBack.count() == 0L) + assert(readBack.schema.map(_.name) == Seq("id", "name")) + } + } + } + + test("a failing task aborts, cleans up its staging file, and the retry succeeds") { + assume(isSpark40Plus, "Requires the WriteFilesExec seam") + // CometWriteFilesExec.executeTask must call committer.abortTask and rethrow. Injecting the + // failure through the commit protocol rather than the data lets the write get as far as + // creating a staging file, so the cleanup is actually observable. + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val sourcePath = new File(dir, "source.parquet").getAbsolutePath + withNativeWriter { + val df = materializeAsCometSource( + (1 to 100).map(i => (i, s"str_$i")).toDF("id", "name").repartition(1), + sourcePath) + + FailingCommitProtocol.reset() + try { + withSQLConf( + SQLConf.FILE_COMMIT_PROTOCOL_CLASS.key -> classOf[FailingCommitProtocol].getName) { + FailingCommitProtocol.failOnCommitTask = true + val e = intercept[Exception] { + df.write.parquet(outputPath) + } + assert( + causeChain(e).exists(_.getMessage == FailingCommitProtocol.message), + s"Expected the injected failure to propagate, got: $e") + assert( + FailingCommitProtocol.abortTaskCalled, + "CometWriteFilesExec must call committer.abortTask when the task fails") + } + } finally { + FailingCommitProtocol.reset() + } + + // The failed job leaves no data behind: no committed part files and no staging tree. + assert( + listPartFileNames(outputPath).isEmpty, + s"A failed write must not leave part files: ${listPartFileNames(outputPath)}") + assert( + !new File(outputPath, "_temporary").exists(), + "A failed write must not leave a _temporary staging directory behind") + + // The same write, without the injected failure, still produces correct output. + val plan = captureWritePlan(p => df.write.mode(SaveMode.Overwrite).parquet(p), outputPath) + assertHasCometNativeWriteExec(plan) + checkAnswer(spark.read.parquet(outputPath), df) + } + } + } + + test("the Spark 3.x opt-in key still enables native writes on Spark 4.0+") { + assume(isSpark40Plus, "Requires the WriteFilesExec seam") + // The opt-in moved from DataWritingCommandExec to WriteFilesExec with the operator. Jobs that + // had already enabled the experimental writer must not silently fall back, so the old key is + // registered as a deprecated alternative - and CometExecRule resolves operator configs by name + // rather than through the ConfigEntry, so isOperatorAllowIncompat has to honor it too. + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val sourcePath = new File(dir, "source.parquet").getAbsolutePath + withSQLConf( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true") { + val df = materializeAsCometSource( + (1 to 100).map(i => (i, s"str_$i")).toDF("id", "name"), + sourcePath) + + withSQLConf(CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true") { + assertHasCometNativeWriteExec( + captureWritePlan(p => df.write.mode(SaveMode.Overwrite).parquet(p), outputPath)) + } + + // An explicitly set new key wins over the deprecated one. + withSQLConf( + CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + CometConf.COMET_OPERATOR_WRITE_FILES_ALLOW_INCOMPAT.key -> "false") { + assertNoCometNativeWriteExec( + captureWritePlan(p => df.write.mode(SaveMode.Overwrite).parquet(p), outputPath)) + } + } + } + } + private def createTestData(inputDir: File): String = { val inputPath = new File(inputDir, "input.parquet").getAbsolutePath val schema = FuzzDataGenerator.generateSchema( @@ -884,7 +1135,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", + nativeWriteAllowIncompatKey -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax")(f) } @@ -963,39 +1214,56 @@ class CometParquetWriterSuite extends CometTestBase { } } + /** + * The operator that carries a native write, which differs by Spark version: on 4.0+ Comet + * replaces only `WriteFilesExec` with [[CometWriteFilesExec]] and leaves Spark's write + * framework in place, while on 3.x it replaces the whole `DataWritingCommandExec` with + * [[CometNativeWriteExec]]. See `CometWriteFiles` / `CometDataWritingCommand`. + */ + private def isNativeWriteExec(plan: SparkPlan): Boolean = plan match { + case _: CometWriteFilesExec => isSpark40Plus + case _: CometNativeWriteExec => !isSpark40Plus + case _ => false + } + + /** The opt-in config key for native writes, which moved with the operator on Spark 4.0+. */ + private def nativeWriteAllowIncompatKey: String = + if (isSpark40Plus) { + CometConf.COMET_OPERATOR_WRITE_FILES_ALLOW_INCOMPAT.key + } else { + CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key + } + private def assertHasCometNativeWriteExec(plan: SparkPlan): Unit = { var nativeWriteCount = 0 - plan.foreach { - case _: CometNativeWriteExec => - nativeWriteCount += 1 - case d: DataWritingCommandExec => - d.child.foreach { - case _: CometNativeWriteExec => - nativeWriteCount += 1 - case _ => - } - case _ => - } + plan.foreach(p => if (isNativeWriteExec(p)) nativeWriteCount += 1) assert( nativeWriteCount == 1, - s"Expected exactly one CometNativeWriteExec in the plan, but found $nativeWriteCount:\n${plan.treeString}") + "Expected exactly one native write operator in the plan, but found " + + s"$nativeWriteCount:\n${plan.treeString}") + + if (isSpark40Plus) { + // On 4.0+ the command is left in the plan on purpose for a fully native write, so it must + // not be reported as a fallback - otherwise extended explain tells users an accelerated + // write was not accelerated, and skews the "Comet accelerated N of M operators" count. + plan.foreach { + case d: DataWritingCommandExec => + val reasons = d.getTagValue(CometExplainInfo.FALLBACK_REASONS).getOrElse(Set.empty) + assert( + reasons.isEmpty, + s"A fully native write must not tag ${d.nodeName} as a fallback, got: $reasons") + case _ => + } + } } private def assertNoCometNativeWriteExec(plan: SparkPlan): Unit = { - val hasNativeWrite = plan.exists { - case _: CometNativeWriteExec => true - case d: DataWritingCommandExec => - d.child.exists { - case _: CometNativeWriteExec => true - case _ => false - } - case _ => false - } + val hasNativeWrite = plan.exists(isNativeWriteExec) assert( !hasNativeWrite, - s"Expected no CometNativeWriteExec in the plan, but found one:\n${plan.treeString}") + s"Expected no native write operator in the plan, but found one:\n${plan.treeString}") } private def writeWithCometNativeWriteExec( @@ -1106,7 +1374,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", + nativeWriteAllowIncompatKey -> "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", @@ -1183,3 +1451,41 @@ class CometParquetWriterSuite extends CometTestBase { } } + +/** + * A commit protocol that fails `commitTask` on demand, to exercise the abort branch of + * `CometWriteFilesExec.executeTask`. + * + * Failing at commit rather than mid-write means the task has already asked for a staging file and + * written it, so `abortTask` has something real to clean up. Spark instantiates this reflectively + * per job via `spark.sql.sources.commitProtocolClass`, hence the companion object for the flags - + * the tests run in `local[*]`, so executors share the driver's JVM and see them. + */ +class FailingCommitProtocol(jobId: String, path: String, dynamicPartitionOverwrite: Boolean) + extends SQLHadoopMapReduceCommitProtocol(jobId, path, dynamicPartitionOverwrite) { + + override def commitTask( + taskContext: TaskAttemptContext): FileCommitProtocol.TaskCommitMessage = { + if (FailingCommitProtocol.failOnCommitTask) { + throw new IOException(FailingCommitProtocol.message) + } + super.commitTask(taskContext) + } + + override def abortTask(taskContext: TaskAttemptContext): Unit = { + FailingCommitProtocol.abortTaskCalled = true + super.abortTask(taskContext) + } +} + +object FailingCommitProtocol { + val message = "injected commitTask failure" + + @volatile var failOnCommitTask: Boolean = false + @volatile var abortTaskCalled: Boolean = false + + def reset(): Unit = { + failOnCommitTask = false + abortTaskCalled = false + } +} 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 1c3500426c3..4679125541a 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 @@ -36,7 +36,6 @@ 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.exchange.ShuffleExchangeLike import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.internal.SQLConf @@ -44,7 +43,7 @@ import org.apache.spark.sql.types.{IntegerType, StructType} import org.apache.spark.unsafe.Platform import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus +import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus} import org.apache.comet.serde.OperatorOuterClass class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -815,8 +814,13 @@ 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", + // The native write opt-in moved from DataWritingCommandExec to WriteFilesExec on + // Spark 4.0+, where Comet replaces only the per-task write. See CometWriteFiles. + (if (isSpark40Plus) { + CometConf.COMET_OPERATOR_WRITE_FILES_ALLOW_INCOMPAT.key + } else { + CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key + }) -> "true", SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax") { spark.sparkContext.setJobGroup(jobGroupId, "native parquet write output metrics") try { From df968953cbe8585983d9821db475723b8be1ba51 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 7 Sep 2026 08:15:12 -0600 Subject: [PATCH 2/9] refactor: take a single deprecated alternative for operator incompat configs Matches the reviewed form on #5293: ConfigBuilder mutates in place, so the Seq destructuring was rebinding the same object. Only one operator has an alternative and there is no reason to expect more. --- .../main/scala/org/apache/comet/CometConf.scala | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index b6bebfe9b09..d515e006bd0 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -1012,7 +1012,7 @@ object CometConf extends ShimCometConf { * itself - otherwise an old key would read `true` from the entry while the planner saw `false`. */ private val operatorIncompatAlternatives = - scala.collection.mutable.Map.empty[String, Seq[String]] + scala.collection.mutable.Map.empty[String, String] /** Spark 3.x only: native writes replace the whole `DataWritingCommandExec` there. */ val COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT: ConfigEntry[Boolean] = @@ -1026,7 +1026,7 @@ object CometConf extends ShimCometConf { val COMET_OPERATOR_WRITE_FILES_ALLOW_INCOMPAT: ConfigEntry[Boolean] = createOperatorIncompatConfig( "WriteFilesExec", - Seq(getOperatorAllowIncompatConfigKey("DataWritingCommandExec"))) + Some(getOperatorAllowIncompatConfigKey("DataWritingCommandExec"))) /** Create a config to enable a specific operator */ private def createExecEnabledConfig( @@ -1053,17 +1053,17 @@ object CometConf extends ShimCometConf { private def createOperatorIncompatConfig( name: String, - alternatives: Seq[String] = Nil): ConfigEntry[Boolean] = { + alternative: Option[String] = None): ConfigEntry[Boolean] = { val configKey = getOperatorAllowIncompatConfigKey(name) val envVar = configKeyToEnvVar(configKey) - // ConfigBuilder mutates in place, so the result of withAlternative is `builder` itself. 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") - if (alternatives.nonEmpty) { - operatorIncompatAlternatives.put(name, alternatives) - builder.withAlternative(alternatives.head, alternatives.tail: _*) + alternative.foreach { alt => + operatorIncompatAlternatives.put(name, alt) + // ConfigBuilder mutates in place, so the result of withAlternative is `builder` itself. + builder.withAlternative(alt) } builder.booleanConf.createWithEnvVarOrDefault(envVar, false) } @@ -1092,7 +1092,7 @@ object CometConf extends ShimCometConf { val value = CometConfDeprecations.readWithAlternatives( conf, getOperatorAllowIncompatConfigKey(name), - operatorIncompatAlternatives.getOrElse(name, Nil)) + operatorIncompatAlternatives.get(name).toSeq) value != null && value.toLowerCase(Locale.ROOT) == "true" } From d36ad9ffa6401dda8fb9465b712105e11306cfee Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 8 Sep 2026 10:25:47 -0600 Subject: [PATCH 3/9] fix: escape the output path before parsing it as a URI The output path reaches both write serdes as `Path.toString`, which decodes percent escapes: a directory containing a space or a literal `%` yields a string that is not a valid URI, and `URI.create` throws on it. On Spark 4.0+ that exception escaped `CometWriteFiles.convert` and failed the query. On 3.x `CometDataWritingCommand.convert` caught it and silently handed the write back to Spark, so the native writer was never used for those paths. Round-trip through `Path` instead, which re-escapes. Only the scheme and authority reach `extractObjectStoreOptions`, but parsing has to succeed to get at them. Also assert that the INSERT INTO visibility test's write actually went native, rather than inferring it from the read-back. --- .../operator/CometDataWritingCommand.scala | 6 ++-- .../serde/operator/CometWriteFiles.scala | 8 +++-- .../parquet/CometParquetWriterSuite.scala | 35 +++++++++++++++++-- 3 files changed, 42 insertions(+), 7 deletions(-) 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 index aec4d71198e..d13d46010bf 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala @@ -19,7 +19,6 @@ package org.apache.comet.serde.operator -import java.net.URI import java.util.Locale import scala.jdk.CollectionConverters._ @@ -129,8 +128,11 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec // Collect S3/cloud storage configurations val session = op.session val hadoopConf = session.sessionState.newHadoopConfWithOptions(cmd.options) + // `outputPath` is `Path.toString`, which is not a valid URI string: it leaves spaces and + // literal `%` unescaped, so `URI.create` would throw (and the catch below would silently + // give the write back to Spark). Going through `Path` escapes them again. val objectStoreOptions = - NativeConfig.extractObjectStoreOptions(hadoopConf, URI.create(outputPath)) + NativeConfig.extractObjectStoreOptions(hadoopConf, cmd.outputPath.toUri) objectStoreOptions.foreach { case (key, value) => writerOpBuilder.putObjectStoreOptions(key, value) } 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 ca81f15b688..18b267cdb0c 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 @@ -19,9 +19,9 @@ package org.apache.comet.serde.operator -import java.net.URI import java.util.Locale +import org.apache.hadoop.fs.Path import org.apache.parquet.hadoop.ParquetOutputFormat import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap import org.apache.spark.sql.comet.{CometNativeExec, CometWriteFilesExec} @@ -141,8 +141,12 @@ object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { // 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) + // The tag holds `Path.toString`, which is not a valid URI string: it leaves spaces and + // literal `%` unescaped, so `URI.create` would throw. Round-tripping through `Path` escapes + // them again. Only the scheme and authority matter to `extractObjectStoreOptions`, but + // parsing has to succeed to get at them. NativeConfig - .extractObjectStoreOptions(hadoopConf, URI.create(outputPath)) + .extractObjectStoreOptions(hadoopConf, new Path(outputPath).toUri) .foreach { case (key, value) => writerOpBuilder.putObjectStoreOptions(key, value) } } 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 97b3f417c29..c37d489ad5f 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -865,6 +865,27 @@ class CometParquetWriterSuite extends CometTestBase { } } + test("output path needing URI escaping still writes natively") { + // The output path reaches the serde as `Path.toString`, which leaves characters that are + // illegal in a URI unescaped. `URI.create` then throws for a space or a literal `%`, costing + // the write its native path - silently on Spark 3.x, where the serde catches the failure and + // hands the write back to Spark. + Seq("dir with space", "dir%with%percent", "dir with space and %25").foreach { dirName => + withTempPath { dir => + val outputPath = new File(new File(dir, dirName), "output.parquet").getAbsolutePath + val sourcePath = new File(dir, "source.parquet").getAbsolutePath + withNativeWriter { + val df = materializeAsCometSource( + (1 to 100).map(i => (i, s"str_$i")).toDF("id", "name"), + sourcePath) + val plan = captureWritePlan(p => df.write.parquet(p), outputPath) + assertHasCometNativeWriteExec(plan) + checkAnswer(spark.read.parquet(outputPath), df) + } + } + } + } + // --------------------------------------------------------------------------------------------- // Spark 4.0+ only. These cover behavior that comes from leaving Spark's write framework in // place, which is only possible where `V1WritesUtils.getWriteFilesOpt` matches the @@ -937,7 +958,11 @@ class CometParquetWriterSuite extends CometTestBase { 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") + // Assert the write itself went native: a fallback to Spark's writer would make the + // read-back pass for the wrong reason. + assertHasCometNativeWriteExec( + captureWritePlan( + 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) } @@ -1172,7 +1197,11 @@ class CometParquetWriterSuite extends CometTestBase { * @return * The captured execution plan */ - private def captureWritePlan(writeOp: String => Unit, outputPath: String): SparkPlan = { + private def captureWritePlan(writeOp: String => Unit, outputPath: String): SparkPlan = + captureWritePlan(writeOp(outputPath)) + + /** As above, for a write that names its own target (an `INSERT INTO`, for example). */ + private def captureWritePlan(writeOp: => Unit): SparkPlan = { var capturedPlan: Option[QueryExecution] = None val listener = new org.apache.spark.sql.util.QueryExecutionListener { @@ -1191,7 +1220,7 @@ class CometParquetWriterSuite extends CometTestBase { spark.listenerManager.register(listener) try { - writeOp(outputPath) + writeOp // Wait for listener to be called with timeout val maxWaitTimeMs = 15000 From 4095ea6fdbd67cbe4aa2fdd2736fed053790e035 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 9 Sep 2026 09:22:59 -0600 Subject: [PATCH 4/9] fix: decline HDFS writes whose output path needs URI escaping --- .../operator/CometDataWritingCommand.scala | 4 +++ .../serde/operator/CometWriteFiles.scala | 4 +++ .../serde/operator/NativeWriteUtils.scala | 31 +++++++++++++++++++ .../parquet/CometParquetWriterSuite.scala | 27 ++++++++++++++++ 4 files changed, 66 insertions(+) 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 index d13d46010bf..ed2ac8b7cd3 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala @@ -63,6 +63,10 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec return Unsupported(Some("Supported output filesystems: local, HDFS")) } + NativeWriteUtils + .escapedHdfsDestination(cmd.outputPath.toString) + .foreach(reason => return Unsupported(Some(reason))) + if (cmd.bucketSpec.isDefined) { return Unsupported(Some("Bucketed writes are not supported")) } 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 18b267cdb0c..900962575f1 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 @@ -82,6 +82,10 @@ object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { return Unsupported(Some("Supported output filesystems: local, HDFS")) } + NativeWriteUtils + .escapedHdfsDestination(outputPath) + .foreach(reason => return Unsupported(Some(reason))) + if (op.bucketSpec.isDefined) { return Unsupported(Some("Bucketed writes are not supported")) } diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala index 8d01806c8e9..a4bf29229cc 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala @@ -19,6 +19,7 @@ package org.apache.comet.serde.operator +import org.apache.hadoop.fs.Path import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.comet.serde.OperatorOuterClass @@ -51,4 +52,34 @@ object NativeWriteUtils { .setScan(scan.build()) .build()) } + + /** + * A fallback reason when `outputPath` is an HDFS destination whose path needs percent-escaping, + * or `None` when the write can proceed. + * + * Comet and Spark disagree about what such a path names. The native side reaches HDFS through + * `create_hdfs_object_store`, which hands `url.path()` -- still escaped -- to + * `object_store::path::Path::parse`, so the native writer creates a directory literally called + * `dir%20with%20space`. Spark's committer, meanwhile, works with the unescaped Hadoop `Path` + * and commits `dir with space`. Job commit then succeeds while the data sits somewhere else, + * which is worse than not accelerating the write. + * + * Local `file:` destinations are unaffected and deliberately not gated here: they go through a + * different object-store constructor that does not retain the escaping. + */ + def escapedHdfsDestination(outputPath: String): Option[String] = { + if (!outputPath.startsWith("hdfs:")) return None + val uri = new Path(outputPath).toUri + // `getPath` decodes, `getRawPath` does not. They differ exactly when the path contains + // something the URI form had to escape. + val raw = uri.getRawPath + val decoded = uri.getPath + if (raw != null && decoded != null && raw != decoded) { + Some( + "HDFS output paths needing URI escaping are not supported: the native writer would " + + s"write to the escaped path while Spark commits the unescaped one ($decoded)") + } else { + None + } + } } 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 c37d489ad5f..2c924224ff6 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -42,6 +42,7 @@ import org.apache.spark.sql.types.{ArrayType, LongType, MapType, Metadata, Metad import org.apache.comet.{CometConf, CometExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} +import org.apache.comet.serde.operator.NativeWriteUtils import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, SchemaGenOptions} class CometParquetWriterSuite extends CometTestBase { @@ -886,6 +887,32 @@ class CometParquetWriterSuite extends CometTestBase { } } + test("HDFS output paths needing URI escaping are declined at planning") { + // The local case above writes natively, but HDFS cannot: `create_hdfs_object_store` hands the + // still-escaped `url.path()` to `object_store::path::Path::parse`, so the native writer would + // create `dir%20with%20space` while Spark's committer commits `dir with space`. Job commit + // would succeed with the data somewhere else, so the write has to stay on Spark until the + // native path handling preserves Hadoop filenames. + Seq( + "hdfs://ns/dir with space/output.parquet", + "hdfs://ns/dir%with%percent/output.parquet", + "hdfs://ns/nested/dir with space/output.parquet").foreach { path => + assert( + NativeWriteUtils.escapedHdfsDestination(path).isDefined, + s"expected $path to be declined") + } + + // Unescaped HDFS paths, and local paths of any shape, are unaffected. + Seq( + "hdfs://ns/plain/output.parquet", + "file:///tmp/dir with space/output.parquet", + "file:///tmp/dir%with%percent/output.parquet").foreach { path => + assert( + NativeWriteUtils.escapedHdfsDestination(path).isEmpty, + s"expected $path to be accepted") + } + } + // --------------------------------------------------------------------------------------------- // Spark 4.0+ only. These cover behavior that comes from leaving Spark's write framework in // place, which is only possible where `V1WritesUtils.getWriteFilesOpt` matches the From 39e02de39908ec4973e7a233cf6f202e86795b0e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 9 Sep 2026 14:07:49 -0600 Subject: [PATCH 5/9] fix: also decline Unicode HDFS write destinations The raw/decoded URI comparison only catches what java.net.URI had to escape, and java.net.URI leaves non-ASCII path characters alone, so an hdfs://ns/cafe/output destination was admitted. percent_encoding's should_percent_encode is !byte.is_ascii() || set.contains(byte), so the native parser escapes every non-ASCII byte regardless of the encode set and the writer creates caf%C3%A9 outside Spark's staging directory. The guard now also declines any character the native parser rewrites. The ASCII half of that set was determined against the locked url 2.5 crate by parsing hdfs://ns/prepost/output for every printable ASCII c: space, ", #, <, >, ?, backtick, { and } are rewritten and the rest survive, so partition directories and Spark's _temporary attempt paths still qualify. The comment no longer claims the Java comparison detects the divergence on its own; both conditions are kept because the Java one still catches a literal % that the native parser leaves alone. Tests add accented (precomposed and combining), CJK, emoji and nested non-ASCII cases plus the remaining escaped ASCII characters, built from code points since scalastyle forbids non-ASCII source. Disabling the new condition makes the accented case fail, so the Java comparison alone demonstrably does not cover it. --- .../serde/operator/NativeWriteUtils.scala | 59 +++++++++++++++---- .../parquet/CometParquetWriterSuite.scala | 40 +++++++++++-- 2 files changed, 83 insertions(+), 16 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala index a4bf29229cc..0819f26b5fa 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala @@ -54,15 +54,48 @@ object NativeWriteUtils { } /** - * A fallback reason when `outputPath` is an HDFS destination whose path needs percent-escaping, - * or `None` when the write can proceed. + * ASCII characters the native URL parser percent-encodes inside a path. Determined against the + * locked `url` 2.5 crate by parsing `hdfs://ns/prepost/output` for every printable ASCII `c` + * and comparing `url.path()` with the input: these nine are rewritten and the rest survive, + * including `%`, `[`, `\`, `]`, `^` and `|`. Control characters and DEL are handled separately + * in [[needsNativeUrlEscaping]] rather than listed here. + */ + private val nativeUrlEscapedAscii: Set[Char] = + Set(' ', '"', '#', '<', '>', '?', '`', '{', '}') + + /** + * Whether the native URL parser would rewrite `path`, so that the name it creates on HDFS + * differs from the Hadoop filename Spark commits. + * + * `percent_encoding`'s `should_percent_encode` is `!byte.is_ascii() || set.contains(byte)`, so + * every non-ASCII byte is escaped regardless of the encode set. That is the case Java's URI + * comparison cannot see, because `java.net.URI` leaves non-ASCII path characters alone: a name + * holding U+00E9 comes back identically from `getRawPath` and `getPath`, while the native + * parser produces `caf%C3%A9`. + */ + private def needsNativeUrlEscaping(path: String): Boolean = + path.exists(c => c < ' ' || c > '~' || nativeUrlEscapedAscii.contains(c)) + + /** + * A fallback reason when `outputPath` is an HDFS destination whose path the native writer and + * Spark would spell differently, or `None` when the write can proceed. + * + * Comet and Spark disagree about what such a path names. The native side receives + * `Path.toString`, which is not a URI string, and reaches HDFS through + * `create_hdfs_object_store`, which hands `url.path()` -- now escaped by the Rust parser -- to + * `object_store::path::Path::parse`. So the native writer creates a directory literally called + * `dir%20with%20space`, or `caf%C3%A9` for a name holding U+00E9. Spark's committer, meanwhile, + * works with the unescaped Hadoop `Path` and commits `dir with space`, or that same U+00E9 name + * unescaped. Job commit then succeeds while the data sits somewhere else, which is worse than + * not accelerating the write. + * + * Two conditions, because neither covers the other: * - * Comet and Spark disagree about what such a path names. The native side reaches HDFS through - * `create_hdfs_object_store`, which hands `url.path()` -- still escaped -- to - * `object_store::path::Path::parse`, so the native writer creates a directory literally called - * `dir%20with%20space`. Spark's committer, meanwhile, works with the unescaped Hadoop `Path` - * and commits `dir with space`. Job commit then succeeds while the data sits somewhere else, - * which is worse than not accelerating the write. + * - the native parser escapes a character, which is the direct statement of the divergence + * and the only condition that catches non-ASCII names; + * - `java.net.URI` had to escape something, which catches a literal `%` in the Hadoop name. + * The native parser leaves `%` alone, so `50%off` reaches `Path::parse` as an invalid + * escape rather than as a rewritten name. * * Local `file:` destinations are unaffected and deliberately not gated here: they go through a * different object-store constructor that does not retain the escaping. @@ -70,14 +103,16 @@ object NativeWriteUtils { def escapedHdfsDestination(outputPath: String): Option[String] = { if (!outputPath.startsWith("hdfs:")) return None val uri = new Path(outputPath).toUri - // `getPath` decodes, `getRawPath` does not. They differ exactly when the path contains - // something the URI form had to escape. val raw = uri.getRawPath val decoded = uri.getPath - if (raw != null && decoded != null && raw != decoded) { + val javaEscaped = raw != null && decoded != null && raw != decoded + // Checked against the string handed to the native writer, which is `outputPath` itself. + val nativeEscaped = needsNativeUrlEscaping(outputPath) + if (javaEscaped || nativeEscaped) { + val shown = if (decoded != null) decoded else outputPath Some( "HDFS output paths needing URI escaping are not supported: the native writer would " + - s"write to the escaped path while Spark commits the unescaped one ($decoded)") + s"write to the escaped path while Spark commits the unescaped one ($shown)") } else { None } 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 2c924224ff6..d491b66348d 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -889,24 +889,56 @@ class CometParquetWriterSuite extends CometTestBase { test("HDFS output paths needing URI escaping are declined at planning") { // The local case above writes natively, but HDFS cannot: `create_hdfs_object_store` hands the - // still-escaped `url.path()` to `object_store::path::Path::parse`, so the native writer would + // now-escaped `url.path()` to `object_store::path::Path::parse`, so the native writer would // create `dir%20with%20space` while Spark's committer commits `dir with space`. Job commit // would succeed with the data somewhere else, so the write has to stay on Spark until the // native path handling preserves Hadoop filenames. + // + // Non-ASCII names are the case a `getRawPath != getPath` comparison alone cannot see: + // `java.net.URI` leaves non-ASCII path characters alone, so both accessors agree, while + // `percent_encoding` escapes every non-ASCII byte regardless of the encode set and the native + // writer creates `caf%C3%A9`. Built from code points because scalastyle forbids non-ASCII + // source characters. + def cp(codePoints: Int*): String = codePoints.map(Character.toChars(_).mkString).mkString + val eAcute = cp(0x00e9) // precomposed LATIN SMALL LETTER E WITH ACUTE + val combiningAcute = cp(0x0301) // COMBINING ACUTE ACCENT, applied to a plain "e" + val cjk = cp(0x65e5, 0x672c, 0x8a9e) // "nihongo" + val emoji = cp(0x1f642) // astral plane, so a surrogate pair on the JVM + val uUmlaut = cp(0x00fc) + Seq( "hdfs://ns/dir with space/output.parquet", "hdfs://ns/dir%with%percent/output.parquet", - "hdfs://ns/nested/dir with space/output.parquet").foreach { path => + "hdfs://ns/nested/dir with space/output.parquet", + s"hdfs://ns/caf$eAcute/output.parquet", + s"hdfs://ns/cafe$combiningAcute/output.parquet", + s"hdfs://ns/$cjk/output.parquet", + s"hdfs://ns/$emoji/output.parquet", + // Non-ASCII in a nested segment rather than the leaf. + s"hdfs://ns/${uUmlaut}ber/nested/output.parquet", + // The remaining ASCII characters the native parser escapes. + "hdfs://ns/quote\"here/output.parquet", + "hdfs://ns/hash#here/output.parquet", + "hdfs://ns/angle/output.parquet", + "hdfs://ns/question?here/output.parquet", + "hdfs://ns/back`tick/output.parquet", + "hdfs://ns/brace{here}/output.parquet").foreach { path => assert( NativeWriteUtils.escapedHdfsDestination(path).isDefined, s"expected $path to be declined") } - // Unescaped HDFS paths, and local paths of any shape, are unaffected. + // Ordinary HDFS paths, and local paths of any shape, are unaffected. Keeping these passing is + // the point of gating on the exact set the native parser rewrites rather than on "not plain + // ASCII alphanumerics", which would decline the partition directories Spark actually writes. Seq( "hdfs://ns/plain/output.parquet", + "hdfs://ns/part-00000-a1b2.c3d4-c000.snappy.parquet", + "hdfs://ns/dt=2026-09-09/hour=17/output.parquet", + "hdfs://ns/_temporary/0/_temporary/attempt_202609091700_0001_m_000000_0/part-0.parquet", "file:///tmp/dir with space/output.parquet", - "file:///tmp/dir%with%percent/output.parquet").foreach { path => + "file:///tmp/dir%with%percent/output.parquet", + s"file:///tmp/caf$eAcute/output.parquet").foreach { path => assert( NativeWriteUtils.escapedHdfsDestination(path).isEmpty, s"expected $path to be accepted") From 268fdb9d2e89139245017133d525dece924ab713 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 13 Sep 2026 08:48:05 -0600 Subject: [PATCH 6/9] fix: address review feedback on the native WriteFilesExec seam Correctness: - Decline HDFS writes whose *file names* would diverge, not just their directory. `mapreduce.output.basename` is caller-controlled and reaches every committed name through `HadoopMapReduceCommitProtocol.getFilename`; a basename holding `?` or `#` makes the native URL parser truncate, so every task writes the same name and they overwrite each other at commit. Adds an execution-time backstop over the complete `newTaskTempFile` path, which a custom commit protocol owns and planning cannot predict. - Read the compression option case-insensitively, as Spark's `ParquetOptions` does. `option("Compression", "lz4_raw")` used to fall through to the SQLConf default, so the unsupported-codec guard was bypassed and Comet wrote SNAPPY into a file Spark had named `.lz4raw.parquet`. The codec is now also re-derived per task from `CodecConfig.from(taskAttemptContext)`, the same place the file extension comes from, so the name and the contents agree by construction. The shared helpers move to `NativeWriteUtils`, which fixes the identical bug on the Spark 3.x path. - Use `Utils.tryWithSafeFinallyAndFailureCallbacks` / `tryWithSafeFinally` in `executeTask` and `writeNatively`, matching `FileFormatWriter`: a failure while aborting or closing the iterator is attached as a suppressed exception instead of replacing the failure that caused it. `statsTrackers` moves inside the guard so a throwing `newTaskInstance` still reaches `abortTask`. Planning and reporting: - Convert `WriteFilesExec` from its enclosing `DataWritingCommandExec` rather than from a separate tag pre-pass, so the output path comes straight from the command that owns it and neither `withNewChildren` copying tags nor "nothing hands us a bare WriteFilesExec" has to hold. - Only skip the fallback reason on `DataWritingCommandExec` when its child really was converted; a write with no native child now says why it fell back. - Drop the node's duplicate `files_written`/`bytes_written`/`rows_written`. `BasicWriteJobStatsTracker` is authoritative here, and the native `bytes_written` reads 0 on HDFS. Tests and docs: - Rust `url_path_rewritten_characters` pins the `url` crate's path encode set, which the JVM guard mirrors; a crate upgrade can no longer reopen the hole with a green build. - New JVM coverage: mixed-case `compression` (honored, and declined when unsupported), the #3426 nested-name INSERT, an empty non-zero partition writing no file, the third-party `WriteTaskStatsTracker` warning, and the basename/committer-path guards. - The abort test no longer claims to demonstrate task retry or speculation. - `installation.md` says which Spark versions its EXPLAIN output applies to. --- docs/source/user-guide/latest/installation.md | 6 +- .../src/execution/operators/parquet_writer.rs | 6 +- native/core/src/parquet/parquet_support.rs | 47 ++++ .../apache/comet/rules/CometExecRule.scala | 54 ++-- .../operator/CometDataWritingCommand.scala | 41 +-- .../serde/operator/CometWriteFiles.scala | 56 ++--- .../serde/operator/NativeWriteUtils.scala | 156 ++++++++++-- .../spark/sql/comet/CometWriteFilesExec.scala | 90 +++++-- .../parquet/CometParquetWriterSuite.scala | 237 +++++++++++++++++- 9 files changed, 547 insertions(+), 146 deletions(-) diff --git a/docs/source/user-guide/latest/installation.md b/docs/source/user-guide/latest/installation.md index 06df8340404..74a6d7dd2b2 100644 --- a/docs/source/user-guide/latest/installation.md +++ b/docs/source/user-guide/latest/installation.md @@ -137,7 +137,7 @@ Create a test Parquet source scala> (0 until 10).toDF("a").write.mode("overwrite").parquet("/tmp/test") ``` -Comet will log output similar to: +Comet will log output similar to this on Spark 4.0 and later: ```shell INFO core/src/lib.rs: Comet native library version $COMET_VERSION initialized @@ -147,6 +147,10 @@ WARN CometExecRule: Comet cannot execute some parts of this plan natively (set s +- LocalTableScan [COMET: Native support for operator LocalTableScanExec is disabled. Set spark.comet.exec.localTableScan.enabled=true to enable it.] ``` +On Spark 3.4 and 3.5 the native writer replaces the whole write command rather than just the +per-task write, so the same message appears on `Execute InsertIntoHadoopFsRelationCommand` and +names `DataWritingCommandExec`. + Query the data from the test source and check: - INFO message shows the native Comet library has been initialized. diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index 187a84fc5cc..46474d7c289 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -514,11 +514,11 @@ impl ExecutionPlan for ParquetWriterExec { Arc::new(Schema::new(fields)) }); - // Spark 4.0+ hands over the exact file to write, chosen by the JVM commit protocol. - // Spark 3.x hands over a working directory instead and expects the writer to name the - // file; that branch goes away with Spark 3.x support. let part_file = match &work_dir { + // Spark 4.0+ hands over the exact file to write, chosen by the JVM commit protocol. None => self.output_path.clone(), + // Spark 3.x hands over a working directory instead and expects the writer to name the + // file; that branch goes away with Spark 3.x support. Some(work_dir) => match task_attempt_id { Some(attempt_id) => format!( "{}/part-{:05}-{:05}.parquet", diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index b601f984d60..6c82401ff14 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -687,6 +687,53 @@ mod tests { prepare_object_store_with_configs(runtime_env, url, &HashMap::new()) } + /// Pins the set of characters the `url` crate rewrites inside a path. + /// + /// `create_hdfs_object_store` hands `url.path()` to `object_store::path::Path::parse`, so a + /// character that the parser escapes - or, for `?` and `#`, treats as a delimiter and + /// truncates at - makes the native writer create a different file from the one Spark's commit + /// protocol chose and later commits. Job commit still succeeds, so the data simply ends up + /// somewhere nobody looks. + /// + /// `NativeWriteUtils.nativeUrlEscapedAscii` on the JVM side keeps a copy of this set and + /// declines such destinations during planning. Nothing else would notice if a `url` upgrade + /// moved a character between the two groups and quietly reopened that hole, so assert the set + /// here rather than in the Scala test, which can only check the copy against itself. + #[test] + fn url_path_rewritten_characters() { + use url::Url; + + let rewritten: String = (' '..='~') + .filter(|c| { + let url = Url::parse(&format!("hdfs://ns/pre{c}post/output")).unwrap(); + url.path() != format!("/pre{c}post/output") + }) + .collect(); + + assert_eq!( + rewritten, " \"#<>?`{}", + "the ASCII characters `url` rewrites inside a path changed; update \ + NativeWriteUtils.nativeUrlEscapedAscii and its tests to match" + ); + + // Every non-ASCII byte is percent-encoded regardless of the encode set. This is the half + // a `java.net.URI` raw/decoded comparison cannot see, because Java leaves non-ASCII path + // characters alone. + for name in [ + "caf\u{e9}", // precomposed e-acute + "cafe\u{301}", // "e" plus a combining acute accent + "\u{65e5}\u{672c}\u{8a9e}", // CJK + "\u{1f642}", // astral plane + ] { + let url = Url::parse(&format!("hdfs://ns/{name}/output")).unwrap(); + assert_ne!( + url.path(), + format!("/{name}/output"), + "expected `url` to percent-encode the non-ASCII name {name}" + ); + } + } + #[cfg(not(feature = "hdfs-opendal"))] #[test] fn test_prepare_object_store() { 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 00246ccce30..a58deaecb5d 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -108,10 +108,12 @@ object CometExecRule { /** * 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 [[org.apache.comet.serde.operator.CometWriteFiles]] learns the target filesystem - and, - * by its absence, that a write comes from some other V1 write command. Only used on Spark 4.0+; - * see `CometWriteFilesExec`. + * `InsertIntoHadoopFsRelationCommand`. `WriteFilesExec` itself has no output path, and + * `CometOperatorSerde` only ever sees the operator, so this is how + * [[org.apache.comet.serde.operator.CometWriteFiles]] learns the target filesystem. Set + * immediately before the one `convertToComet` call that reads it, from the command that owns + * the path; the absence of the tag means the write came from somewhere else and must be + * declined. Only used on Spark 4.0+; see `CometWriteFilesExec`. */ val WRITE_OUTPUT_PATH: TreeNodeTag[String] = TreeNodeTag[String]("comet.writeOutputPath") @@ -402,8 +404,16 @@ case class CometExecRule(session: SparkSession) // therefore Spark's commit protocol, stats trackers and SaveMode handling - in place. // `V1WritesUtils.getWriteFilesOpt` matches the `WriteFilesExecBase` trait there, which is // what lets a Comet node stand in for the write node. See CometWriteFilesExec. - case w: WriteFilesExec if isSpark40Plus => - convertToComet(w, CometWriteFiles).getOrElse(w) + // + // Matched at the command rather than at the write node so that the output path, which + // `WriteFilesExec` does not carry, comes straight from the command that owns it. Converting + // the child from here also means a `WriteFilesExec` without its enclosing command - which + // nothing produces today - is simply left on Spark instead of being converted against a + // stale or missing tag. + case d @ DataWritingCommandExec(cmd: InsertIntoHadoopFsRelationCommand, w: WriteFilesExec) + if isSpark40Plus => + w.setTagValue(CometExecRule.WRITE_OUTPUT_PATH, cmd.outputPath.toString) + d.withNewChildren(Seq(convertToComet(w, CometWriteFiles).getOrElse(w))) // Spark 3.x: `getWriteFilesOpt` matches the concrete `WriteFilesExec` case class, so a // Comet node can never stand in for the write node. Native writes instead replace the whole @@ -509,17 +519,20 @@ case class CometExecRule(session: SparkSession) // Some execs should never be replaced. We include // these cases specially here so we do not add a misleading 'info' message. op - case _: WriteFilesExec if !isSpark40Plus => - // On Spark 3.x the write is converted at the DataWritingCommandExec above, which - // unwraps WriteFilesExec inside convertToComet. Tagging it 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. On - // 4.0+ the node was offered to CometWriteFiles above and already carries a reason. + case _: WriteFilesExec => + // The write is converted at the enclosing DataWritingCommandExec above: on Spark 3.x + // by replacing the whole command, on 4.0+ by converting this child from there. + // Tagging it 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; where the write really did fall back on 4.0+, the node + // already carries the reason CometWriteFiles gave. op - case _: DataWritingCommandExec if isSpark40Plus => + case d: DataWritingCommandExec + if isSpark40Plus && d.child.isInstanceOf[CometWriteFilesExec] => // On Spark 4.0+ DataWritingCommandExec is deliberately left in the plan even for a // fully native write - Comet replaces only its WriteFilesExec child - so tagging it - // would report an accelerated write as a fallback. + // would report an accelerated write as a fallback. A write whose child was not + // converted still falls through to the default case below and gets a reason. op case _ => // The operator was not converted to a Comet plan and no serde handler claimed it, so @@ -541,19 +554,6 @@ 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. Only the Spark 4.0+ path consults this tag. - if (isSpark40Plus) { - 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 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 index ed2ac8b7cd3..01a90f8c8c6 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala @@ -19,11 +19,8 @@ package org.apache.comet.serde.operator -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 @@ -43,9 +40,6 @@ import org.apache.comet.serde.OperatorOuterClass.Operator */ 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) @@ -64,7 +58,9 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec } NativeWriteUtils - .escapedHdfsDestination(cmd.outputPath.toString) + // This writer names its own files `part--.parquet`, so the + // prefix is fixed rather than read from `mapreduce.output.basename`. + .escapedHdfsDestination(cmd.outputPath.toString, "part") .foreach(reason => return Unsupported(Some(reason))) if (cmd.bucketSpec.isDefined) { @@ -75,8 +71,8 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec return Unsupported(Some("Partitioned writes are not supported")) } - val codec = parseCompressionCodec(cmd) - if (!supportedCompressionCodes.contains(codec)) { + val codec = NativeWriteUtils.parseCompressionCodec(cmd.options) + if (!NativeWriteUtils.supportedCompressionCodecs.contains(codec)) { return Unsupported(Some(s"Unsupported compression codec: $codec")) } @@ -106,14 +102,11 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec 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") + val plannedCodec = NativeWriteUtils.parseCompressionCodec(cmd.options) + val codec = NativeWriteUtils.protoCompressionCodec(plannedCodec) match { + case Some(codec) => codec + case None => + withFallbackReason(op, s"Unsupported compression codec: $plannedCodec") return None } @@ -200,18 +193,4 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec 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 index 900962575f1..fffa5d7dc4e 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 @@ -19,10 +19,8 @@ package org.apache.comet.serde.operator -import java.util.Locale - +import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.Path -import org.apache.parquet.hadoop.ParquetOutputFormat import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap import org.apache.spark.sql.comet.{CometNativeExec, CometWriteFilesExec} import org.apache.spark.sql.execution.datasources.WriteFilesExec @@ -43,9 +41,6 @@ import org.apache.comet.serde.OperatorOuterClass.Operator */ 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) @@ -83,7 +78,7 @@ object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { } NativeWriteUtils - .escapedHdfsDestination(outputPath) + .escapedHdfsDestination(outputPath, fileNamePrefix(hadoopConf(op))) .foreach(reason => return Unsupported(Some(reason))) if (op.bucketSpec.isDefined) { @@ -102,8 +97,8 @@ object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { Some("Writes with spark.sql.files.maxRecordsPerFile set are not supported")) } - val codec = parseCompressionCodec(op) - if (!supportedCompressionCodecs.contains(codec)) { + val codec = NativeWriteUtils.parseCompressionCodec(op.options) + if (!NativeWriteUtils.supportedCompressionCodecs.contains(codec)) { return Unsupported(Some(s"Unsupported compression codec: $codec")) } @@ -124,14 +119,13 @@ object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { return None } - 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") + // Planning-time value only, so that a plan can be inspected without a task context. + // CometWriteFilesExec replaces it per task with the codec Parquet names the file after. + val plannedCodec = NativeWriteUtils.parseCompressionCodec(op.options) + val codec = NativeWriteUtils.protoCompressionCodec(plannedCodec) match { + case Some(codec) => codec + case None => + withFallbackReason(op, s"Unsupported compression codec: $plannedCodec") return None } @@ -144,13 +138,12 @@ object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { // 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) // The tag holds `Path.toString`, which is not a valid URI string: it leaves spaces and // literal `%` unescaped, so `URI.create` would throw. Round-tripping through `Path` escapes // them again. Only the scheme and authority matter to `extractObjectStoreOptions`, but // parsing has to succeed to get at them. NativeConfig - .extractObjectStoreOptions(hadoopConf, new Path(outputPath).toUri) + .extractObjectStoreOptions(hadoopConf(op), new Path(outputPath).toUri) .foreach { case (key, value) => writerOpBuilder.putObjectStoreOptions(key, value) } } @@ -170,6 +163,17 @@ object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { private def outputPathOf(op: WriteFilesExec): Option[String] = op.getTagValue(CometExecRule.WRITE_OUTPUT_PATH) + private def hadoopConf(op: WriteFilesExec): Configuration = + op.session.sessionState.newHadoopConfWithOptions(op.options) + + /** + * The leading component of every file name this write will produce. Spark's + * `HadoopMapReduceCommitProtocol.getFilename` reads it from the task configuration, so a write + * option or a session-level Hadoop setting can replace the usual `part`. + */ + private def fileNamePrefix(hadoopConf: Configuration): String = + hadoopConf.get(NativeWriteUtils.BASE_OUTPUT_NAME, NativeWriteUtils.DEFAULT_BASE_OUTPUT_NAME) + /** * Whether Spark would roll to a new file every N rows within a task. * @@ -188,18 +192,4 @@ object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] { .getOrElse(SQLConf.get.maxRecordsPerFile) maxRecordsPerFile > 0 } - - 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/comet/serde/operator/NativeWriteUtils.scala b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala index 0819f26b5fa..fe7edda1fb3 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala @@ -19,8 +19,13 @@ package org.apache.comet.serde.operator +import java.util.Locale + import org.apache.hadoop.fs.Path +import org.apache.parquet.hadoop.ParquetOutputFormat import org.apache.spark.sql.catalyst.plans.QueryPlan +import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap +import org.apache.spark.sql.internal.SQLConf import org.apache.comet.serde.OperatorOuterClass import org.apache.comet.serde.QueryPlanSerde.serializeDataType @@ -31,6 +36,16 @@ import org.apache.comet.serde.QueryPlanSerde.serializeDataType */ object NativeWriteUtils { + /** + * Hadoop's `FileOutputFormat.BASE_OUTPUT_NAME`, which is `protected` there. Spelled out for the + * same reason Spark spells it out in `HadoopMapReduceCommitProtocol.getFilename`, which is + * where a write's effective value is read. + */ + val BASE_OUTPUT_NAME: String = "mapreduce.output.basename" + + /** The file-name prefix a write uses when [[BASE_OUTPUT_NAME]] is not set. */ + val DEFAULT_BASE_OUTPUT_NAME: String = "part" + /** * Build a synthetic `Scan` operator that lets a native write op consume Arrow batches shipped * from the JVM iterator over `plan`'s `executeColumnar()` RDD. @@ -54,18 +69,24 @@ object NativeWriteUtils { } /** - * ASCII characters the native URL parser percent-encodes inside a path. Determined against the - * locked `url` 2.5 crate by parsing `hdfs://ns/prepost/output` for every printable ASCII `c` - * and comparing `url.path()` with the input: these nine are rewritten and the rest survive, + * ASCII characters the native URL parser rewrites inside a path. Determined against the locked + * `url` 2.5 crate by parsing `hdfs://ns/prepost/output` for every printable ASCII `c` and + * comparing `url.path()` with the input: these nine are rewritten and the rest survive, * including `%`, `[`, `\`, `]`, `^` and `|`. Control characters and DEL are handled separately * in [[needsNativeUrlEscaping]] rather than listed here. + * + * `native/core/src/parquet/parquet_support.rs` has a `url_path_rewritten_characters` test that + * fails if a `url` upgrade changes this set, so the two cannot drift apart silently. + * + * `?` and `#` are the worst of the nine: they are not escaped but treated as delimiters, so the + * native path is *truncated* there rather than merely spelled differently. */ private val nativeUrlEscapedAscii: Set[Char] = Set(' ', '"', '#', '<', '>', '?', '`', '{', '}') /** - * Whether the native URL parser would rewrite `path`, so that the name it creates on HDFS - * differs from the Hadoop filename Spark commits. + * Whether the native URL parser would rewrite `s`, so that the name it creates on HDFS differs + * from the Hadoop filename Spark commits. * * `percent_encoding`'s `should_percent_encode` is `!byte.is_ascii() || set.contains(byte)`, so * every non-ASCII byte is escaped regardless of the encode set. That is the case Java's URI @@ -73,15 +94,13 @@ object NativeWriteUtils { * holding U+00E9 comes back identically from `getRawPath` and `getPath`, while the native * parser produces `caf%C3%A9`. */ - private def needsNativeUrlEscaping(path: String): Boolean = - path.exists(c => c < ' ' || c > '~' || nativeUrlEscapedAscii.contains(c)) + private def needsNativeUrlEscaping(s: String): Boolean = + s.exists(c => c < ' ' || c > '~' || nativeUrlEscapedAscii.contains(c)) /** - * A fallback reason when `outputPath` is an HDFS destination whose path the native writer and - * Spark would spell differently, or `None` when the write can proceed. + * Whether the native writer and Spark would spell `path` differently, and how. * - * Comet and Spark disagree about what such a path names. The native side receives - * `Path.toString`, which is not a URI string, and reaches HDFS through + * The native side receives `Path.toString`, which is not a URI string, and reaches HDFS through * `create_hdfs_object_store`, which hands `url.path()` -- now escaped by the Rust parser -- to * `object_store::path::Path::parse`. So the native writer creates a directory literally called * `dir%20with%20space`, or `caf%C3%A9` for a name holding U+00E9. Spark's committer, meanwhile, @@ -98,23 +117,116 @@ object NativeWriteUtils { * escape rather than as a rewritten name. * * Local `file:` destinations are unaffected and deliberately not gated here: they go through a - * different object-store constructor that does not retain the escaping. + * different object-store constructor that keeps the string verbatim. */ - def escapedHdfsDestination(outputPath: String): Option[String] = { - if (!outputPath.startsWith("hdfs:")) return None - val uri = new Path(outputPath).toUri + private def hdfsPathDivergence(path: String): Option[String] = { + if (!path.startsWith("hdfs:")) return None + val uri = new Path(path).toUri val raw = uri.getRawPath val decoded = uri.getPath val javaEscaped = raw != null && decoded != null && raw != decoded - // Checked against the string handed to the native writer, which is `outputPath` itself. - val nativeEscaped = needsNativeUrlEscaping(outputPath) - if (javaEscaped || nativeEscaped) { - val shown = if (decoded != null) decoded else outputPath - Some( - "HDFS output paths needing URI escaping are not supported: the native writer would " + - s"write to the escaped path while Spark commits the unescaped one ($shown)") + // Checked against the string handed to the native writer, which is `path` itself. + if (javaEscaped || needsNativeUrlEscaping(path)) { + Some(if (decoded != null) decoded else path) } else { None } } + + /** + * A fallback reason when a write to `outputPath` would land somewhere the committer is not + * looking, or `None` when the write can proceed. + * + * Two things go into every committed file name, and the native writer has to reproduce both + * byte for byte (see [[hdfsPathDivergence]] for why it may not): + * + * - the destination directory, and + * - `fileNamePrefix`, the basename every file name is built from. On Spark 4.0+ that is + * `mapreduce.output.basename`, which `HadoopMapReduceCommitProtocol.getFilename` + * interpolates into `--`; on 3.x Comet names the files itself and + * the basename is always the literal `part`. A basename holding `?` or `#` is the dangerous + * one: the native URL parser truncates there, so *every* task writes a file with the same + * truncated name and they overwrite each other during commit. + */ + def escapedHdfsDestination(outputPath: String, fileNamePrefix: String): Option[String] = { + if (!outputPath.startsWith("hdfs:")) return None + hdfsPathDivergence(outputPath) + .map(shown => + "HDFS output paths needing URI escaping are not supported: the native writer would " + + s"write to the escaped path while Spark commits the unescaped one ($shown)") + .orElse { + if (needsNativeUrlEscaping(fileNamePrefix)) { + Some( + s"HDFS output file names needing URI escaping are not supported: " + + s"$BASE_OUTPUT_NAME=$fileNamePrefix would make the native writer create a " + + "different file from the one Spark commits") + } else { + None + } + } + } + + /** + * Fail the task if the native writer would not write to exactly `filePath`. + * + * [[escapedHdfsDestination]] declines the shapes Comet can predict at planning time, but the + * path a write actually uses comes from `FileCommitProtocol.newTaskTempFile`, and a custom + * commit protocol can return anything. This is the backstop: it runs before the writer opens + * anything, so the task fails with a clear message rather than committing successfully with the + * data left somewhere else. + */ + def checkNativeWriteDestination(filePath: String): Unit = + hdfsPathDivergence(filePath).foreach { shown => + throw new UnsupportedOperationException( + s"Comet's native Parquet writer cannot write to '$filePath': the path it would create " + + s"on HDFS is not the one the commit protocol chose ($shown). Set " + + "spark.comet.parquet.write.enabled=false to write this table with Spark.") + } + + /** Compression codecs Comet's native Parquet writer can produce. */ + val supportedCompressionCodecs: Set[String] = + Set("none", "uncompressed", "snappy", "lz4", "zstd", "gzip") + + /** + * The compression codec a write will use, resolved exactly as Spark's own `ParquetOptions` + * does: `compression`, then `parquet.compression` (i.e. `ParquetOutputFormat.COMPRESSION`), + * then `spark.sql.parquet.compression.codec`. + * + * The `CaseInsensitiveMap` is not cosmetic. Spark reaches these options through one + * (`ParquetOptions.compressionCodecClassName`) and `DataFrameWriter` passes the caller's keys + * through verbatim, so `option("Compression", "gzip")` is a gzip write to Spark. Reading it + * case-sensitively would both name the file `...-c000.gz.parquet` while writing SNAPPY, and let + * an unsupported codec slip past [[supportedCompressionCodecs]] by falling through to the + * SQLConf default. + */ + def parseCompressionCodec(options: Map[String, String]): String = { + val caseInsensitive = CaseInsensitiveMap(options) + caseInsensitive + .get("compression") + .orElse(caseInsensitive.get(ParquetOutputFormat.COMPRESSION)) + .getOrElse( + SQLConf.get.getConfString( + SQLConf.PARQUET_COMPRESSION.key, + SQLConf.PARQUET_COMPRESSION.defaultValueString)) + .toLowerCase(Locale.ROOT) + } + + /** + * The proto codec for a Spark codec name, or `None` if Comet cannot produce it. + * + * At execution time the name to pass here is `CodecConfig.from(context).getCodec.name()`: + * `ParquetUtils.prepareWrite` resolves the write's options into `parquet.compression` on the + * job configuration, and the file extension comes from that same `CodecConfig`. Reading the + * codec back from there is what makes the file's name and its contents agree by construction + * rather than by two copies of the precedence rule agreeing. + */ + def protoCompressionCodec(codec: String): Option[OperatorOuterClass.CompressionCodec] = + codec.toLowerCase(Locale.ROOT) match { + case "snappy" => Some(OperatorOuterClass.CompressionCodec.Snappy) + case "lz4" => Some(OperatorOuterClass.CompressionCodec.Lz4) + case "zstd" => Some(OperatorOuterClass.CompressionCodec.Zstd) + case "gzip" => Some(OperatorOuterClass.CompressionCodec.Gzip) + case "none" | "uncompressed" => Some(OperatorOuterClass.CompressionCodec.None) + case _ => None + } } 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 1beeb0c8537..c5b0bc44ba6 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 @@ -25,6 +25,7 @@ import scala.jdk.CollectionConverters._ import org.apache.hadoop.mapreduce.{TaskAttemptContext, TaskAttemptID, TaskID, TaskType} import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl +import org.apache.parquet.hadoop.codec.CodecConfig import org.apache.spark.TaskContext import org.apache.spark.internal.Logging import org.apache.spark.internal.io.{FileCommitProtocol, FileNameSpec, SparkHadoopWriterUtils} @@ -35,7 +36,7 @@ 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.execution.metric.SQLMetric import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType import org.apache.spark.sql.vectorized.ColumnarBatch @@ -43,7 +44,7 @@ import org.apache.spark.util.Utils import org.apache.comet.serde.OperatorOuterClass import org.apache.comet.serde.OperatorOuterClass.Operator -import org.apache.comet.serde.operator.schema2Proto +import org.apache.comet.serde.operator.{schema2Proto, NativeWriteUtils} import org.apache.comet.shims.ShimCometWriteFilesExec /** @@ -87,10 +88,14 @@ case class CometWriteFilesExec( 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")) + /** + * No metrics of its own. On this path `BasicWriteJobStatsTracker` is authoritative for the file + * count, byte count and row count, and reports them on the enclosing `Execute + * InsertIntoHadoopFsRelationCommand` node. Republishing the native writer's own counters here + * would only add a second, less trustworthy copy: its `bytes_written` comes from + * `std::fs::metadata`, which returns 0 for a file on HDFS. + */ + override lazy val metrics: Map[String, SQLMetric] = Map.empty override def serializedPlanOpt: SerializedPlan = SerializedPlan(Some(CometExec.serializeNativePlan(nativeOp))) @@ -101,6 +106,14 @@ case class CometWriteFilesExec( * 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. + * + * Note that no `EliminateRedundantTransitions` arm is needed to keep a `ColumnarToRowExec` off + * this node, unlike `CometNativeWriteExec` on the Spark 3.x path. `CometExecRule` is a + * `preColumnarTransitions` rule, so `ApplyColumnarRulesAndInsertTransitions` would normally + * insert one above a columnar child - but for a `V1WriteCommand` under `plannedWriteEnabled` it + * passes `outputsColumnar = write.child.supportsColumnar` instead (`Columnar.scala`), which + * leaves this node alone. That is load-bearing: a transition here would not fail at planning + * but at execution, as a "has write support mismatch". */ override def doExecute(): RDD[InternalRow] = throw new UnsupportedOperationException(s"$nodeName does not support doExecute") @@ -131,9 +144,11 @@ case class CometWriteFilesExec( // 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. - // The write's target schema, not the query output's: 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, nullability and Parquet field IDs. + // The write's target schema rather than the child's output. Spark's analyzer normally makes + // the two agree - `castAndRenameQueryOutput` aliases an INSERT's select list to the target's + // column names and casts nested structs to the target's field names - but `dataColumns` is + // what Spark itself treats as authoritative, and it is the one that stays correct when + // partitioned writes arrive, since `FileFormatWriter` excludes the partition columns from it. val dataSchema = CometUtils.fromAttributes(description.dataColumns) val taskWrite = NativeWriteTask( @@ -197,9 +212,17 @@ object CometWriteFilesExec extends Logging { taskCtx.taskAttemptId().toInt & Integer.MAX_VALUE) committer.setupTask(taskAttemptContext) - val statsTrackers = description.statsTrackers.map(_.newTaskInstance()) - try { + // Same guard FileFormatWriter.executeTask uses. A plain try/catch/finally would let a failure + // while aborting or cleaning up replace the failure that caused the abort; this keeps the + // original and attaches the rest as suppressed exceptions, and marks the task failed so + // Spark's failure listeners run. + Utils.tryWithSafeFinallyAndFailureCallbacks(block = { + // Inside the guard so that a tracker whose newTaskInstance throws still reaches abortTask. + // Spark gets this for free: its trackers are built by the FileFormatDataWriter constructor, + // which runs inside this same block. + val statsTrackers = description.statsTrackers.map(_.newTaskInstance()) + // 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. @@ -209,15 +232,31 @@ object CometWriteFilesExec extends Logging { // naming. The file counter is always 0 until file rolling is supported. val filePath = committer.newTaskTempFile(taskAttemptContext, None, FileNameSpec("", "-c000" + ext)) + // CometWriteFiles declines the destinations whose native spelling it can predict, but the + // path itself comes from the commit protocol, which is replaceable. Check the real thing + // before anything is created: a write that lands outside the committer's staging tree + // still commits successfully, so failing here is the only way the user finds out. + NativeWriteUtils.checkNativeWriteDestination(filePath) + + // `ext` above is the codec's extension, so take the codec from the same place Parquet + // took it - `ParquetUtils.prepareWrite` resolved the write's options into the job + // configuration - rather than from the planning-time value. A file whose name says gzip + // and whose footer says snappy is worse than a write that never happened. + val codec = CodecConfig.from(taskAttemptContext).getCodec + val protoCodec = NativeWriteUtils + .protoCompressionCodec(codec.name()) + .getOrElse(throw new UnsupportedOperationException( + s"Comet's native Parquet writer cannot write $codec")) statsTrackers.foreach(_.newFile(filePath)) - val rowsWritten = writeNatively(taskWrite, filePath, batches, sparkPartitionId) + val rowsWritten = + writeNatively(taskWrite, filePath, protoCodec, 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()) + // `hasNext` already ran the child to completion, so there is nothing to drain or release + // here. Spark's EmptyDirectoryDataWriter writes nothing in this case either. "no file" } @@ -234,12 +273,10 @@ object CometWriteFilesExec extends Logging { // 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 - } + })(catchBlock = { + committer.abortTask(taskAttemptContext) + logError(s"Task ${taskAttemptContext.getTaskAttemptID} aborted") + }) } /** @@ -251,10 +288,12 @@ object CometWriteFilesExec extends Logging { private def writeNatively( taskWrite: NativeWriteTask, filePath: String, + codec: OperatorOuterClass.CompressionCodec, batches: Iterator[ColumnarBatch], partitionId: Int): Long = { val parquetWriter = taskWrite.nativeOp.getParquetWriter.toBuilder .setOutputPath(filePath) + .setCompression(codec) .clearColumnNames() .addAllColumnNames(taskWrite.dataColumnNames.asJava) .clearOutputSchema() @@ -276,12 +315,14 @@ object CometWriteFilesExec extends Logging { broadcastedHadoopConfForEncryption = None, encryptedFilePaths = Seq.empty) - try { + // `close()` propagates teardown failures, including from the final native metrics update, so + // a plain `finally` would let a cleanup error replace the write error that caused it. + Utils.tryWithSafeFinally { // The native writer emits no batches; draining performs the write. while (execIterator.hasNext) { execIterator.next().close() } - } finally { + } { execIterator.close() } @@ -300,8 +341,11 @@ object CometWriteFilesExec extends Logging { * 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. + * + * Visible for testing: nothing Spark ships lets a caller install a third-party tracker on a V1 + * write, so this is the only way to exercise the warning. */ - private def recordRows( + def recordRows( statsTrackers: Seq[WriteTaskStatsTracker], filePath: String, count: Long): Unit = { 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 d491b66348d..57c9bc97ce1 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -21,22 +21,25 @@ package org.apache.comet.parquet import java.io.{File, IOException} +import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ import scala.util.{Random, Using} import org.apache.hadoop.fs.{FileSystem, Path} import org.apache.hadoop.mapreduce.TaskAttemptContext +import org.apache.logging.log4j.Level import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.metadata.CompressionCodecName import org.apache.parquet.hadoop.util.HadoopInputFile import org.apache.parquet.schema.{MessageType, Type} import org.apache.spark.internal.io.FileCommitProtocol import org.apache.spark.sql.{AnalysisException, CometTestBase, DataFrame, Row, SaveMode} +import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometNativeWriteExec, 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.SQLHadoopMapReduceCommitProtocol -import org.apache.spark.sql.functions.{array, map, struct, when} +import org.apache.spark.sql.execution.datasources.{BasicWriteTaskStats, SQLHadoopMapReduceCommitProtocol, WriteTaskStats, WriteTaskStatsTracker} +import org.apache.spark.sql.functions.{array, col, map, struct, when} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, LongType, MapType, Metadata, MetadataBuilder, StringType, StructField, StructType} @@ -431,6 +434,66 @@ class CometParquetWriterSuite extends CometTestBase { } } + test("parquet write honors a mixed-case compression option") { + // Spark resolves write options through a CaseInsensitiveMap (ParquetOptions) and + // DataFrameWriter hands the caller's keys through verbatim, so `Compression` really does ask + // for gzip. Reading it case-sensitively would fall through to the SQLConf default and write + // SNAPPY into a file Spark had already named `...-c000.gz.parquet`. + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val df = spark.range(0, 100).selectExpr("id", "cast(id as string) as name") + + withSQLConf( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + nativeWriteAllowIncompatKey -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + SQLConf.PARQUET_COMPRESSION.key -> "snappy") { + + val plan = captureWritePlan( + path => df.write.option("Compression", "gzip").parquet(path), + outputPath) + assertHasCometNativeWriteExec(plan) + } + + checkAnswer(spark.read.parquet(outputPath), df.collect()) + assertParquetCodec(outputPath, CompressionCodecName.GZIP) + if (isSpark40Plus) { + // Spark names the file; Comet fills it. The extension is the only externally visible + // statement of the codec, so it has to agree with the footer. (On 3.x the native writer + // invents a name with no codec suffix, so there is nothing to compare.) + listPartFileNames(outputPath).foreach { name => + assert(name.endsWith(".gz.parquet"), s"Expected a gzip file name, got '$name'") + } + } + } + } + + test("parquet write with a mixed-case unsupported compression codec falls back to Spark") { + assume(isSpark35Plus, "lz4_raw was added in Spark 3.5") + // The other half of the case bug, and the dangerous half: a case-sensitive read misses the + // option entirely, finds `snappy` in the SQLConf, decides the codec is supported, and writes + // SNAPPY bytes into a file Spark named `...-c000.lz4raw.parquet`. + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val df = spark.range(0, 100).selectExpr("id", "cast(id as string) as name") + + withSQLConf( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + nativeWriteAllowIncompatKey -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + SQLConf.PARQUET_COMPRESSION.key -> "snappy") { + + val plan = captureWritePlan( + path => df.write.option("Compression", "lz4_raw").parquet(path), + outputPath) + assertNoCometNativeWriteExec(plan) + } + + checkAnswer(spark.read.parquet(outputPath), df.collect()) + assertParquetCodec(outputPath, CompressionCodecName.LZ4_RAW) + } + } + test("parquet write with unsupported compression codec falls back to Spark") { assume(isSpark35Plus, "lz4_raw was added in Spark 3.5") withTempPath { dir => @@ -887,7 +950,7 @@ class CometParquetWriterSuite extends CometTestBase { } } - test("HDFS output paths needing URI escaping are declined at planning") { + test("HDFS destinations needing URI escaping are declined at planning") { // The local case above writes natively, but HDFS cannot: `create_hdfs_object_store` hands the // now-escaped `url.path()` to `object_store::path::Path::parse`, so the native writer would // create `dir%20with%20space` while Spark's committer commits `dir with space`. Job commit @@ -924,7 +987,7 @@ class CometParquetWriterSuite extends CometTestBase { "hdfs://ns/back`tick/output.parquet", "hdfs://ns/brace{here}/output.parquet").foreach { path => assert( - NativeWriteUtils.escapedHdfsDestination(path).isDefined, + NativeWriteUtils.escapedHdfsDestination(path, "part").isDefined, s"expected $path to be declined") } @@ -940,9 +1003,42 @@ class CometParquetWriterSuite extends CometTestBase { "file:///tmp/dir%with%percent/output.parquet", s"file:///tmp/caf$eAcute/output.parquet").foreach { path => assert( - NativeWriteUtils.escapedHdfsDestination(path).isEmpty, + NativeWriteUtils.escapedHdfsDestination(path, "part").isEmpty, s"expected $path to be accepted") } + + // The directory is only half of the committed path. `mapreduce.output.basename` puts + // caller-controlled text into every file name, and `?`/`#` are worse there than anywhere + // else: the native URL parser treats them as delimiters and truncates, so every task would + // write the same file name and they would overwrite each other during commit. + val plainHdfs = "hdfs://ns/plain/output.parquet" + Seq("part?x", "part#x", "part with space", s"caf$eAcute").foreach { basename => + assert( + NativeWriteUtils.escapedHdfsDestination(plainHdfs, basename).isDefined, + s"expected basename '$basename' to be declined") + } + Seq("part", "out", "data_v2", "part-of-it").foreach { basename => + assert( + NativeWriteUtils.escapedHdfsDestination(plainHdfs, basename).isEmpty, + s"expected basename '$basename' to be accepted") + } + assert( + NativeWriteUtils + .escapedHdfsDestination("file:///tmp/plain/output.parquet", "part?x") + .isEmpty, + "local writes use the path verbatim, so the basename cannot diverge there") + + // Planning can only decline what it can predict. The path a task actually writes comes from + // FileCommitProtocol.newTaskTempFile, which a custom commit protocol owns, so the same check + // runs again at execution time - as a hard failure, because by then the only alternative is + // committing successfully with the data somewhere else. + NativeWriteUtils.checkNativeWriteDestination( + "hdfs://ns/out/_temporary/0/attempt_1_m_0_0/part-00000-abc-c000.snappy.parquet") + NativeWriteUtils.checkNativeWriteDestination("file:/tmp/out/part?x-00000.parquet") + intercept[UnsupportedOperationException] { + NativeWriteUtils.checkNativeWriteDestination( + "hdfs://ns/out/_temporary/0/attempt_1_m_0_0/part?x-00000-abc-c000.snappy.parquet") + } } // --------------------------------------------------------------------------------------------- @@ -1118,11 +1214,123 @@ class CometParquetWriterSuite extends CometTestBase { } } - test("a failing task aborts, cleans up its staging file, and the retry succeeds") { + test("INSERT INTO ... SELECT writes the target table's column names") { + assume(isSpark40Plus, "Requires the WriteFilesExec seam") + // https://github.com/apache/datafusion-comet/issues/3426, which is Spark's own + // `INSERT INTO TABLE - complex type but different names` (sql/core InsertSuite) with a + // top-level rename added. On the Spark 3.x writer this returns no rows at all (#3521), so + // the scenario is a regression test for the whole design rather than for one line: Comet + // takes the schema from `WriteJobDescription.dataColumns`, but Spark's + // `castAndRenameQueryOutput` has already aliased the query's output to the target's names + // and cast the struct to the target's nested names, so reading the child's output instead + // would produce the same file today. `dataColumns` is still the right source - it is what + // Spark guarantees, and it is the only one that stays correct once partitioned writes are + // supported, since those exclude the partition columns from the data schema. + withTempPath { dir => + val targetPath = new File(dir, "target").getAbsolutePath + withTable("comet_rename_target", "comet_rename_source") { + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sql( + "CREATE TABLE comet_rename_source(id bigint, s struct) " + + "USING parquet") + sql( + "CREATE TABLE comet_rename_target(total bigint, p struct) " + + s"USING parquet LOCATION '$targetPath'") + sql("INSERT INTO comet_rename_source SELECT 1, named_struct('a', 'x', 'b', 'y')") + } + + withNativeWriter { + assertHasCometNativeWriteExec( + captureWritePlan( + sql("INSERT INTO comet_rename_target SELECT id + 1, s FROM comet_rename_source"))) + } + + // Read the names out of the file itself: the catalog would report the target's schema + // whatever the file said, which is exactly how this went unnoticed. + assertParquetSchemas(targetPath) { schema => + assert( + schema.getFields.asScala.map(_.getName) == Seq("total", "p"), + s"Expected the target table's column names in the written file, got $schema") + val nested = schema.getFields.asScala.last.asGroupType() + assert( + nested.getFields.asScala.map(_.getName) == Seq("c", "d"), + s"Expected the target table's nested field names in the written file, got $schema") + } + checkAnswer(spark.table("comet_rename_target"), Row(2L, Row("x", "y")) :: Nil) + } + } + } + + test("an empty partition writes no file and still commits") { + assume(isSpark40Plus, "Requires the WriteFilesExec seam") + // executeTask's `sparkPartitionId != 0 && !batches.hasNext` branch must skip newTaskTempFile + // altogether and still commit the task, matching FileFormatWriter's EmptyDirectoryDataWriter. + // Hash-partitioning into eight and keeping a single id leaves at most one partition with + // rows, so at most two files can appear: that one, plus partition 0's schema-only file when + // the row did not land there. + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val sourcePath = new File(dir, "source.parquet").getAbsolutePath + withNativeWriter { + // AQE would coalesce the eight partitions back down to one on data this small, which + // would remove the empty partitions the test is about. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val df = materializeAsCometSource( + (1 to 100).map(i => (i, s"str_$i")).toDF("id", "name"), + sourcePath) + .repartition(8, col("id")) + .where("id = 7") + + val plan = captureWritePlan(p => df.write.parquet(p), outputPath) + assertHasCometNativeWriteExec(plan) + + val partFiles = listPartFileNames(outputPath) + assert(partFiles.nonEmpty, "The partition holding the row must have written a file") + assert( + partFiles.size <= 2, + s"Empty partitions must not write files, but got ${partFiles.size}: $partFiles") + checkAnswer(spark.read.parquet(outputPath), Row(7, "str_7") :: Nil) + } + } + } + } + + test("a third-party WriteTaskStatsTracker is warned that it gets counts, not row contents") { + // The documented known limitation. `WriteTaskStatsTracker.newRow` is a per-row callback and + // Comet has columnar batches, so rather than materialize every row to hand it straight back + // it passes an empty one - exactly right for BasicWriteTaskStatsTracker, which ignores the + // row, and wrong for anything that inspects it. Nothing Spark ships lets a V1 write install a + // third-party tracker, so drive recordRows directly. + val tracker = new RecordingStatsTracker + val appender = new LogAppender("third-party WriteTaskStatsTracker warning") + withLogAppender(appender, Seq(classOf[CometWriteFilesExec].getName), Some(Level.WARN)) { + CometWriteFilesExec.recordRows(Seq(tracker), "/tmp/part-00000.parquet", 3) + } + + assert(tracker.rows.size == 3, s"Expected 3 row callbacks, got ${tracker.rows.size}") + assert( + tracker.rows.forall(_._1 == "/tmp/part-00000.parquet"), + "Every callback must name the file being written") + assert( + tracker.rows.forall(_._2.numFields == 0), + "The known limitation is that the row is empty, so assert it rather than assume it") + assert( + appender.loggingEvents.exists( + _.getMessage.getFormattedMessage.contains(classOf[RecordingStatsTracker].getName)), + "A tracker that is not BasicWriteTaskStatsTracker must be warned by name, got: " + + appender.loggingEvents.map(_.getMessage.getFormattedMessage).mkString("\n")) + } + + test("a failing task aborts and cleans up its staging file") { assume(isSpark40Plus, "Requires the WriteFilesExec seam") // CometWriteFilesExec.executeTask must call committer.abortTask and rethrow. Injecting the // failure through the commit protocol rather than the data lets the write get as far as // creating a staging file, so the cleanup is actually observable. + // + // The second write at the end is a fresh write of the same data, not an automatic task retry + // and not a speculative attempt: it only shows that the failed job left nothing behind that + // would break the next one. Task-attempt isolation itself is Spark's `newTaskTempFile`, and + // is still untested here. withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath val sourcePath = new File(dir, "source.parquet").getAbsolutePath @@ -1577,3 +1785,20 @@ object FailingCommitProtocol { abortTaskCalled = false } } + +/** + * A `WriteTaskStatsTracker` that is not Spark's own, recording what Comet hands it. + * + * Stands in for a third-party tracker, which Comet cannot supply with row contents. See + * `CometWriteFilesExec.recordRows`. + */ +class RecordingStatsTracker extends WriteTaskStatsTracker { + val rows: ArrayBuffer[(String, InternalRow)] = ArrayBuffer.empty + + override def newPartition(partitionValues: InternalRow): Unit = {} + override def newFile(filePath: String): Unit = {} + override def closeFile(filePath: String): Unit = {} + override def newRow(filePath: String, row: InternalRow): Unit = rows += ((filePath, row)) + override def getFinalStats(taskCommitTime: Long): WriteTaskStats = + BasicWriteTaskStats(Seq.empty, 0, 0, rows.size) +} From 5f5397f5049b1bfc6fe27e6c8e926ea2d0bee959 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 13 Sep 2026 08:58:39 -0600 Subject: [PATCH 7/9] fix: drop a redundant string interpolator flagged by scalafix RedundantSyntax --- .../org/apache/comet/serde/operator/NativeWriteUtils.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala index fe7edda1fb3..504ab17492a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala @@ -157,7 +157,7 @@ object NativeWriteUtils { .orElse { if (needsNativeUrlEscaping(fileNamePrefix)) { Some( - s"HDFS output file names needing URI escaping are not supported: " + + "HDFS output file names needing URI escaping are not supported: " + s"$BASE_OUTPUT_NAME=$fileNamePrefix would make the native writer create a " + "different file from the one Spark commits") } else { From 6970255d0ac5b90b478a3645986ead41910c3afe Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 18 Sep 2026 17:26:27 -0600 Subject: [PATCH 8/9] fix: decline percent-bearing HDFS basenames at planning mapreduce.output.basename=part%foo passed the planning guard and then failed checkNativeWriteDestination at execution, aborting the job instead of falling back to Spark's writer. needsNativeUrlEscaping deliberately excludes % because the url crate leaves it alone, so only Java's raw-vs-decoded comparison sees it. The basename now goes through the same hdfsPathDivergence predicate the task guard uses, applied to the path the basename produces. That closes the gap for % and for the other characters java.net.URI escapes but the url crate keeps ([, ^, |), and makes the two guards agree by construction rather than by two character sets staying in sync. Also updates CometEmptyRelationParquetWriterSuite, which main added while this branch was open: a native empty relation is the zero-partition input CometWriteFilesExec swaps a single-partition RDD in for, so on 4.0+ that write is accelerated rather than declined. --- .../serde/operator/NativeWriteUtils.scala | 19 ++-- .../spark/sql/comet/CometWriteFilesExec.scala | 11 +- .../parquet/CometParquetWriterSuite.scala | 105 +++++++++++++----- .../parquet/CometParquetWriterTestBase.scala | 31 +++++- ...CometEmptyRelationParquetWriterSuite.scala | 12 +- 5 files changed, 131 insertions(+), 47 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala index 504ab17492a..d8247e99a7f 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala @@ -147,6 +147,13 @@ object NativeWriteUtils { * the basename is always the literal `part`. A basename holding `?` or `#` is the dangerous * one: the native URL parser truncates there, so *every* task writes a file with the same * truncated name and they overwrite each other during commit. + * + * The basename is checked by running [[hdfsPathDivergence]] over the path it produces rather + * than over the name alone. Everything else `getFilename` interpolates -- the split number, the + * job id and the codec extension -- is Comet-independent ASCII, so the probe below covers the + * whole committed file name. Sharing the one predicate with [[checkNativeWriteDestination]] is + * also what keeps planning from admitting a basename the task guard then aborts on: a literal + * `%` is invisible to [[needsNativeUrlEscaping]] but not to `java.net.URI`. */ def escapedHdfsDestination(outputPath: String, fileNamePrefix: String): Option[String] = { if (!outputPath.startsWith("hdfs:")) return None @@ -155,14 +162,10 @@ object NativeWriteUtils { "HDFS output paths needing URI escaping are not supported: the native writer would " + s"write to the escaped path while Spark commits the unescaped one ($shown)") .orElse { - if (needsNativeUrlEscaping(fileNamePrefix)) { - Some( - "HDFS output file names needing URI escaping are not supported: " + - s"$BASE_OUTPUT_NAME=$fileNamePrefix would make the native writer create a " + - "different file from the one Spark commits") - } else { - None - } + hdfsPathDivergence(s"$outputPath/$fileNamePrefix").map(_ => + "HDFS output file names needing URI escaping are not supported: " + + s"$BASE_OUTPUT_NAME=$fileNamePrefix would make the native writer create a " + + "different file from the one Spark commits") } } 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 c5b0bc44ba6..10b0b9db15b 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 @@ -127,11 +127,12 @@ case class CometWriteFilesExec( val childRDD = child.executeColumnar() - // SPARK-23271 (defensive): a zero-partition input would spawn no task and therefore write no - // file at all, so the output directory would carry no schema for readers. Spark's own - // WriteFilesExec swaps in a dummy single-partition RDD for exactly this case. In practice - // CometWriteFiles.requiresNativeChildren rules out the sources (LocalTableScan) that produce - // a zero-partition RDD, but the swap is kept to match Spark's semantics if that ever changes. + // SPARK-23271: a zero-partition input would spawn no task and therefore write no file at all, + // so the output directory would carry no schema for readers. Spark's own WriteFilesExec swaps + // in a dummy single-partition RDD for exactly this case. AQE reaches it by collapsing a + // completed empty shuffle into a CometEmptyRelationExec, which is why + // CometDataWritingCommand declines empty-relation inputs on the Spark 3.x path (#5303): its + // writer has nowhere to put this swap. See CometEmptyRelationParquetWriterSuite. val writeRDD = if (childRDD.getNumPartitions == 0) { sparkContext.parallelize(Seq.empty[ColumnarBatch], 1) } else { 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 99af76c3a6e..9062eb94c3f 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -1012,12 +1012,28 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { // else: the native URL parser treats them as delimiters and truncates, so every task would // write the same file name and they would overwrite each other during commit. val plainHdfs = "hdfs://ns/plain/output.parquet" - Seq("part?x", "part#x", "part with space", s"caf$eAcute").foreach { basename => + val declinedBasenames = Seq( + "part?x", + "part#x", + "part with space", + s"caf$eAcute", + // A literal `%` is the case the native-escaping set alone cannot see: the `url` parser + // leaves `%` untouched, so only Java's escaping catches it. `part%25` is the same + // character arriving as an escape sequence rather than as a stray sign. + "part%foo", + "part%25", + // Escaped by `java.net.URI` but not by the native parser, so likewise only the Java half + // of the check sees them. + "part[0]", + "part^x", + "part|x") + declinedBasenames.foreach { basename => assert( NativeWriteUtils.escapedHdfsDestination(plainHdfs, basename).isDefined, s"expected basename '$basename' to be declined") } - Seq("part", "out", "data_v2", "part-of-it").foreach { basename => + val acceptedBasenames = Seq("part", "out", "data_v2", "part-of-it") + acceptedBasenames.foreach { basename => assert( NativeWriteUtils.escapedHdfsDestination(plainHdfs, basename).isEmpty, s"expected basename '$basename' to be accepted") @@ -1039,6 +1055,29 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { NativeWriteUtils.checkNativeWriteDestination( "hdfs://ns/out/_temporary/0/attempt_1_m_0_0/part?x-00000-abc-c000.snappy.parquet") } + + // The two guards have to agree exactly. Planning admitting something the task guard then + // rejects is not a safe direction to be wrong in: the write has already been accepted, so + // the job aborts at execution time instead of quietly falling back to Spark's writer. Check + // that on the name `HadoopMapReduceCommitProtocol.getFilename` actually builds. + def committedPath(basename: String): String = + s"$plainHdfs/_temporary/0/_temporary/attempt_202609091700_0001_m_000000_0/" + + s"$basename-00000-a1b2c3d4-e5f6-c000.snappy.parquet" + (declinedBasenames ++ acceptedBasenames).foreach { basename => + val declinedAtPlanning = + NativeWriteUtils.escapedHdfsDestination(plainHdfs, basename).isDefined + val rejectedAtRuntime = + try { + NativeWriteUtils.checkNativeWriteDestination(committedPath(basename)) + false + } catch { + case _: UnsupportedOperationException => true + } + assert( + declinedAtPlanning == rejectedAtRuntime, + s"basename '$basename': declined at planning = $declinedAtPlanning, but the task guard " + + s"on ${committedPath(basename)} says $rejectedAtRuntime") + } } // --------------------------------------------------------------------------------------------- @@ -1100,6 +1139,37 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + test("a custom output basename is honored on local storage") { + assume(isSpark40Plus, "Requires the WriteFilesExec seam") + // `escapedHdfsDestination` gates the basename on HDFS only, where the native URL parser would + // rename the file out from under the committer. Local writes hand the path to the native + // writer verbatim, so this is the control that keeps that guard from being widened: `part%foo` + // is declined on HDFS and has to keep working here. + Seq("out", "part%foo").foreach { basename => + 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.option(NativeWriteUtils.BASE_OUTPUT_NAME, basename).parquet(p), + outputPath) + assertHasCometNativeWriteExec(plan) + } + val written = + new File(outputPath).listFiles().map(_.getName).filter(_.endsWith(".parquet")) + assert( + written.nonEmpty && written.forall(_.startsWith(s"$basename-")), + s"expected every data file to be named '$basename-...', found: " + + written.mkString(", ")) + checkAnswer(spark.read.parquet(outputPath), df) + } + } + } + } + test("INSERT INTO ... SELECT is visible to subsequent reads") { assume(isSpark40Plus, "Requires the WriteFilesExec seam") // https://github.com/apache/datafusion-comet/issues/3521 - reads returned no rows because the @@ -1190,10 +1260,9 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { // of the output must see the write's schema, not fail. Comet reaches this in two ways - if the // native child has one partition producing no batches, the partition-0 branch of executeTask // writes a metadata-only file; if it produces zero partitions, doExecuteWrite swaps in a dummy - // single-partition RDD to get to the same branch. This test exercises the reachable path - // (filtered Comet scan yielding an empty batch iterator); the zero-partition swap is defensive - // because CometWriteFiles.requiresNativeChildren rules out the sources (LocalTableScan) that - // would otherwise produce a zero-partition RDD. + // single-partition RDD to get to the same branch. This test exercises the first; the + // zero-partition swap is reached by an AQE-collapsed empty relation and is covered by + // CometEmptyRelationParquetWriterSuite. withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath val sourcePath = new File(dir, "source.parquet").getAbsolutePath @@ -1446,30 +1515,6 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } - private def assertHasCometNativeWriteExec(plan: SparkPlan): Unit = { - var nativeWriteCount = 0 - plan.foreach(p => if (isNativeWriteExec(p)) nativeWriteCount += 1) - - assert( - nativeWriteCount == 1, - "Expected exactly one native write operator in the plan, but found " + - s"$nativeWriteCount:\n${plan.treeString}") - - if (isSpark40Plus) { - // On 4.0+ 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 writeWithCometNativeWriteExec( inputPath: String, outputPath: String, diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala index 3d84e7299f5..f3e59e75101 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala @@ -22,9 +22,10 @@ package org.apache.comet.parquet import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.comet.{CometNativeWriteExec, CometWriteFilesExec} import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.command.DataWritingCommandExec import org.apache.spark.sql.internal.SQLConf -import org.apache.comet.CometConf +import org.apache.comet.{CometConf, CometExplainInfo} import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus abstract class CometParquetWriterTestBase extends CometTestBase { @@ -103,8 +104,8 @@ abstract class CometParquetWriterTestBase extends CometTestBase { /** * The operator that carries a native write, which differs by Spark version: on 4.0+ Comet - * replaces only `WriteFilesExec` with [[CometWriteFilesExec]] and leaves Spark's write framework - * in place, while on 3.x it replaces the whole `DataWritingCommandExec` with + * replaces only `WriteFilesExec` with [[CometWriteFilesExec]] and leaves Spark's write + * framework in place, while on 3.x it replaces the whole `DataWritingCommandExec` with * [[CometNativeWriteExec]]. See `CometWriteFiles` / `CometDataWritingCommand`. */ protected def isNativeWriteExec(plan: SparkPlan): Boolean = plan match { @@ -113,6 +114,30 @@ abstract class CometParquetWriterTestBase extends CometTestBase { case _ => false } + protected def assertHasCometNativeWriteExec(plan: SparkPlan): Unit = { + var nativeWriteCount = 0 + plan.foreach(p => if (isNativeWriteExec(p)) nativeWriteCount += 1) + + assert( + nativeWriteCount == 1, + "Expected exactly one native write operator in the plan, but found " + + s"$nativeWriteCount:\n${plan.treeString}") + + if (isSpark40Plus) { + // On 4.0+ 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 _ => + } + } + } + protected def assertNoCometNativeWriteExec(plan: SparkPlan): Unit = { val hasNativeWrite = plan.exists(isNativeWriteExec) diff --git a/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala b/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala index b5580f80ed0..18f9ee2ec0a 100644 --- a/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala +++ b/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala @@ -62,9 +62,19 @@ class CometEmptyRelationParquetWriterSuite extends CometParquetWriterTestBase { assert(readback.schema == StructType(empty.schema.map(_.copy(nullable = true)))) assert(readback.collect().isEmpty) } - assertNoCometNativeWriteExec(plan) if (nativeEmpty) { assert(collect(plan) { case e: CometEmptyRelationExec => e }.nonEmpty) + // A native empty relation is a zero-partition RDD, which is what + // CometDataWritingCommand declines on the Spark 3.x path: that writer only maps + // existing partitions, so no task runs and no file is written. Spark 4.0+ goes + // through the WriteFilesExec seam instead, where CometWriteFilesExec swaps in a + // dummy single-partition RDD exactly as Spark's own WriteFilesExec does, so + // partition 0 still writes the schema-only file the readback above needs. + assertHasCometNativeWriteExec(plan) + } else { + // Spark's own EmptyRelationExec is not a Comet operator, so + // CometWriteFiles.requiresNativeChildren keeps the write on Spark. + assertNoCometNativeWriteExec(plan) } } } From a69b76c130aaacc7a6d2bab274b8e24f81711d9e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 18 Sep 2026 17:51:47 -0600 Subject: [PATCH 9/9] fix: drop imports left unused by moving the write assertions into the base --- .../org/apache/comet/parquet/CometParquetWriterSuite.scala | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) 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 9062eb94c3f..c46524e2917 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -35,15 +35,14 @@ import org.apache.parquet.schema.{MessageType, Type} import org.apache.spark.internal.io.FileCommitProtocol import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SaveMode} import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometNativeWriteExec, CometScanExec, CometWriteFilesExec} +import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometScanExec, CometWriteFilesExec} import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan} -import org.apache.spark.sql.execution.command.DataWritingCommandExec import org.apache.spark.sql.execution.datasources.{BasicWriteTaskStats, SQLHadoopMapReduceCommitProtocol, WriteTaskStats, WriteTaskStatsTracker} import org.apache.spark.sql.functions.{array, col, map, struct, when} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, LongType, MapType, Metadata, MetadataBuilder, StringType, StructField, StructType} -import org.apache.comet.{CometConf, CometExplainInfo} +import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} import org.apache.comet.serde.operator.NativeWriteUtils import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, SchemaGenOptions}