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 42b2800463a..c32f35fcf81 100644 --- a/docs/source/user-guide/latest/installation.md +++ b/docs/source/user-guide/latest/installation.md @@ -181,16 +181,20 @@ 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 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.] ``` +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/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 6a6ea5aa4d9..2182e7de584 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -117,9 +117,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 23ffddef0a8..9897ed626c6 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) + 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", + 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 38b940475e4..8f030da455b 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/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index 1964f531746..4683161a4b1 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -1277,6 +1277,53 @@ mod tests { ); } + /// 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/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index aaa1af076ea..6a6284fb4c1 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -897,11 +897,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 b4455f63c4c..d2369a13c27 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -1024,9 +1024,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, 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", + Some(getOperatorAllowIncompatConfigKey("DataWritingCommandExec"))) + /** Create a config to enable a specific operator */ private def createExecEnabledConfig( exec: String, @@ -1050,15 +1070,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, + alternative: Option[String] = None): ConfigEntry[Boolean] = { val configKey = getOperatorAllowIncompatConfigKey(name) val envVar = configKeyToEnvVar(configKey) - conf(configKey) + val builder = conf(configKey) .category(CATEGORY_EXEC) .doc(s"Whether to allow incompatibility for operator: $name. " + s"False by default. Can be overridden with $envVar env variable") - .booleanConf - .createWithEnvVarOrDefault(envVar, false) + 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) } def isExprEnabled(name: String, conf: SQLConf = SQLConf.get): Boolean = { @@ -1082,7 +1108,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.get(name).toSeq) + 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 12463f6f139..265c84a0dfe 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -38,7 +38,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 @@ -109,6 +109,17 @@ 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, 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") + /** * Tag set on a `ShuffleExchangeExec` that should be left as a plain Spark shuffle rather than * wrapped in `CometShuffleExchangeExec`. See `tagRedundantColumnarShuffle`. @@ -392,15 +403,35 @@ 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. + // + // 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 + // 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 @@ -487,14 +518,24 @@ 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 => + // 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 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. 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 @@ -805,7 +846,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/CometDataWritingCommand.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala index c717e28d86c..6af2cc18ba6 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,12 +19,8 @@ package org.apache.comet.serde.operator -import java.net.URI -import java.util.Locale - import scala.jdk.CollectionConverters._ -import org.apache.parquet.hadoop.ParquetOutputFormat import org.apache.spark.SparkException import org.apache.spark.sql.comet.{CometEmptyRelationExec, CometNativeExec, CometNativeWriteExec, CometScanWrapper} import org.apache.spark.sql.execution.SparkPlan @@ -47,9 +43,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) @@ -77,6 +70,12 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec return Unsupported(Some("Supported output filesystems: local, HDFS")) } + NativeWriteUtils + // 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) { return Unsupported(Some("Bucketed writes are not supported")) } @@ -85,8 +84,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")) } @@ -124,14 +123,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 } @@ -150,8 +146,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) } @@ -215,18 +214,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 new file mode 100644 index 00000000000..fffa5d7dc4e --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometWriteFiles.scala @@ -0,0 +1,195 @@ +/* + * 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 org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path +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] { + + 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")) + } + + NativeWriteUtils + .escapedHdfsDestination(outputPath, fileNamePrefix(hadoopConf(op))) + .foreach(reason => return Unsupported(Some(reason))) + + 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 = NativeWriteUtils.parseCompressionCodec(op.options) + if (!NativeWriteUtils.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 + } + + // 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 + } + + // `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 => + // 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(op), new Path(outputPath).toUri) + .foreach { case (key, value) => writerOpBuilder.putObjectStoreOptions(key, value) } + } + + Some( + Operator + .newBuilder() + .setPlanId(op.id) + .addChildren(scanOperator) + .setParquetWriter(writerOpBuilder.build()) + .build()) + } + + override def createExec(nativeOp: Operator, op: WriteFilesExec): CometNativeExec = + CometWriteFilesExec(nativeOp, originalPlan = op, child = op.child) + + /** The write's output path, recorded on the node by `CometExecRule`. */ + private def outputPathOf(op: WriteFilesExec): Option[String] = + op.getTagValue(CometExecRule.WRITE_OUTPUT_PATH) + + private def 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. + * + * `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 + } +} 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..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 @@ -19,7 +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 @@ -30,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. @@ -51,4 +67,169 @@ object NativeWriteUtils { .setScan(scan.build()) .build()) } + + /** + * 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 `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 + * 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(s: String): Boolean = + s.exists(c => c < ' ' || c > '~' || nativeUrlEscapedAscii.contains(c)) + + /** + * Whether the native writer and Spark would spell `path` differently, and how. + * + * 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: + * + * - 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 keeps the string verbatim. + */ + 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 `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. + * + * 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 + 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 { + 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") + } + } + + /** + * 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 new file mode 100644 index 00000000000..10b0b9db15b --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala @@ -0,0 +1,388 @@ +/* + * 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.parquet.hadoop.codec.CodecConfig +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 +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, NativeWriteUtils} +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" + + /** + * 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))) + + 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. + * + * 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") + + 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: 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 { + 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 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( + 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) + + // 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. + 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)) + // 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, protoCodec, batches, sparkPartitionId) + recordRows(statsTrackers, filePath, rowsWritten) + statsTrackers.foreach(_.closeFile(filePath)) + filePath + } else { + // `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" + } + + 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))))) + })(catchBlock = { + committer.abortTask(taskAttemptContext) + logError(s"Task ${taskAttemptContext.getTaskAttemptID} aborted") + }) + } + + /** + * 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, + 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() + .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) + + // `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() + } + } { + 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. + * + * 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. + */ + 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 75520a035d9..c46524e2917 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,32 @@ package org.apache.comet.parquet -import java.io.File +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, DataFrame, Row, SaveMode} -import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometNativeWriteExec, CometScanExec} +import org.apache.spark.sql.catalyst.InternalRow +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.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} import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus +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 CometParquetWriterTestBase { @@ -77,7 +83,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { 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 +105,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { 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 +139,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { 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 +368,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { 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 +389,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { 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 +414,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { 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") { @@ -428,6 +433,66 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + 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 => @@ -436,7 +501,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { 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 +928,552 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + 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) + } + } + } + } + + 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 + // 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", + 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, "part").isDefined, + s"expected $path to be declined") + } + + // 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", + s"file:///tmp/caf$eAcute/output.parquet").foreach { path => + assert( + 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" + 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") + } + 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") + } + 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") + } + + // 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") + } + } + + // --------------------------------------------------------------------------------------------- + // 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("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 + // 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 { + // 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) + } + } + + 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 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 + 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("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 + 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( @@ -903,25 +1514,6 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } - 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 _ => - } - - assert( - nativeWriteCount == 1, - s"Expected exactly one CometNativeWriteExec in the plan, but found $nativeWriteCount:\n${plan.treeString}") - } - private def writeWithCometNativeWriteExec( inputPath: String, outputPath: String, @@ -1030,7 +1622,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { 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", @@ -1107,3 +1699,58 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + +/** + * 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 + } +} + +/** + * 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) +} 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 96b543e3761..f3e59e75101 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala @@ -20,23 +20,32 @@ package org.apache.comet.parquet import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.comet.CometNativeWriteExec +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 { protected 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) } + /** The opt-in config key for native writes, which moved with the operator on Spark 4.0+. */ + protected def nativeWriteAllowIncompatKey: String = + if (isSpark40Plus) { + CometConf.COMET_OPERATOR_WRITE_FILES_ALLOW_INCOMPAT.key + } else { + CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key + } + /** * Captures the execution plan during a write operation. * @@ -47,7 +56,11 @@ abstract class CometParquetWriterTestBase extends CometTestBase { * @return * The captured execution plan */ - protected def captureWritePlan(writeOp: String => Unit, outputPath: String): SparkPlan = { + protected 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). */ + protected def captureWritePlan(writeOp: => Unit): SparkPlan = { var capturedPlan: Option[QueryExecution] = None val listener = new org.apache.spark.sql.util.QueryExecutionListener { @@ -66,7 +79,7 @@ abstract class CometParquetWriterTestBase extends CometTestBase { spark.listenerManager.register(listener) try { - writeOp(outputPath) + writeOp // Wait for listener to be called with timeout val maxWaitTimeMs = 15000 @@ -89,19 +102,47 @@ abstract class CometParquetWriterTestBase extends CometTestBase { } } - protected 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 + /** + * 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`. + */ + protected def isNativeWriteExec(plan: SparkPlan): Boolean = plan match { + case _: CometWriteFilesExec => isSpark40Plus + case _: CometNativeWriteExec => !isSpark40Plus + 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) 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}") } } 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 { 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) } } }