peterxcli commented on code in PR #5763:
URL: https://github.com/apache/datafusion-comet/pull/5763#discussion_r3997233974
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -513,6 +541,19 @@ case class CometExecRule(session: SparkSession)
}
}
+ // `WriteFilesExec` does not carry the write's output path, but
CometWriteFiles needs it to
+ // decide whether the target filesystem is supported. Record it from the
enclosing command
+ // before the bottom-up walk reaches the write node. The absence of the
tag also tells
+ // CometWriteFiles that the write is not an
InsertIntoHadoopFsRelationCommand and must be
+ // declined. Only the Spark 4.0+ path consults this tag.
+ if (isSpark40Plus) {
+ plan.foreach {
+ case DataWritingCommandExec(cmd: InsertIntoHadoopFsRelationCommand, w:
WriteFilesExec) =>
+ w.setTagValue(CometExecRule.WRITE_OUTPUT_PATH,
cmd.outputPath.toString)
Review Comment:
Is it possible that we match the enclosing `DataWritingCommandExec` and pass
`cmd.outputPath` explicitly when converting its WriteFilesExec child, avoiding
this separate tag set and get?
##########
native/core/src/execution/operators/parquet_writer.rs:
##########
@@ -510,15 +514,18 @@ impl ExecutionPlan for ParquetWriterExec {
Arc::new(Schema::new(fields))
});
- // Generate part file name for this partition
- // If using FileCommitProtocol (work_dir is set), include
task_attempt_id in the filename
- let part_file = if let Some(attempt_id) = task_attempt_id {
- format!(
- "{}/part-{:05}-{:05}.parquet",
- work_dir, self.partition_id, attempt_id
- )
- } else {
- format!("{}/part-{:05}.parquet", work_dir, self.partition_id)
+ // Spark 4.0+ hands over the exact file to write, chosen by the JVM
commit protocol.
+ // Spark 3.x hands over a working directory instead and expects the
writer to name the
+ // file; that branch goes away with Spark 3.x support.
+ let part_file = match &work_dir {
+ None => self.output_path.clone(),
+ Some(work_dir) => match task_attempt_id {
+ Some(attempt_id) => format!(
+ "{}/part-{:05}-{:05}.parquet",
+ work_dir, self.partition_id, attempt_id
+ ),
+ None => format!("{}/part-{:05}.parquet", work_dir,
self.partition_id),
+ },
Review Comment:
nit, want to separate the comment into branches.
```suggestion
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),
},
```
##########
spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala:
##########
@@ -863,6 +866,338 @@ class CometParquetWriterSuite extends CometTestBase {
}
}
+ test("output path needing URI escaping still writes natively") {
+ // The output path reaches the serde as `Path.toString`, which leaves
characters that are
+ // illegal in a URI unescaped. `URI.create` then throws for a space or a
literal `%`, costing
+ // the write its native path - silently on Spark 3.x, where the serde
catches the failure and
+ // hands the write back to Spark.
+ Seq("dir with space", "dir%with%percent", "dir with space and
%25").foreach { dirName =>
+ withTempPath { dir =>
+ val outputPath = new File(new File(dir, dirName),
"output.parquet").getAbsolutePath
+ val sourcePath = new File(dir, "source.parquet").getAbsolutePath
+ withNativeWriter {
+ val df = materializeAsCometSource(
+ (1 to 100).map(i => (i, s"str_$i")).toDF("id", "name"),
+ sourcePath)
+ val plan = captureWritePlan(p => df.write.parquet(p), outputPath)
+ assertHasCometNativeWriteExec(plan)
+ checkAnswer(spark.read.parquet(outputPath), df)
+ }
+ }
+ }
+ }
+
+ test("HDFS output paths needing URI escaping are declined at planning") {
+ // The local case above writes natively, but HDFS cannot:
`create_hdfs_object_store` hands the
+ // 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<here>/output.parquet",
+ "hdfs://ns/question?here/output.parquet",
+ "hdfs://ns/back`tick/output.parquet",
+ "hdfs://ns/brace{here}/output.parquet").foreach { path =>
+ assert(
+ NativeWriteUtils.escapedHdfsDestination(path).isDefined,
+ s"expected $path to be declined")
+ }
+
+ // 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).isEmpty,
+ s"expected $path to be accepted")
+ }
+ }
+
+ //
---------------------------------------------------------------------------------------------
+ // Spark 4.0+ only. These cover behavior that comes from leaving Spark's
write framework in
+ // place, which is only possible where `V1WritesUtils.getWriteFilesOpt`
matches the
+ // `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-<partition>-<uuid>-c<counter>.<codec>.parquet. Committers that
track individual files
+ // and tools that parse these names depend on it.
+ withTempPath { dir =>
+ val outputPath = new File(dir, "output.parquet").getAbsolutePath
+ withTempPath { srcDir =>
+ val df = materializeAsCometSource(
+ (1 to 100).map(i => (i, s"n_$i")).toDF("id", "name"),
+ new File(srcDir, "src.parquet").getAbsolutePath)
+ withNativeWriter {
+ withSQLConf(SQLConf.PARQUET_COMPRESSION.key -> "snappy") {
+ val plan = captureWritePlan(p => df.write.parquet(p), outputPath)
+ assertHasCometNativeWriteExec(plan)
+ }
+ }
+ }
+
+ val partFiles = listPartFileNames(outputPath)
+ assert(partFiles.nonEmpty, s"No part files written to $outputPath")
+ val namePattern =
"""part-\d{5}-[0-9a-f\-]{36}-c\d{3}\.snappy\.parquet""".r
+ partFiles.foreach { name =>
+ assert(
+ namePattern.pattern.matcher(name).matches(),
+ s"File name '$name' does not match Spark's part-file naming
convention")
+ }
+ }
+ }
+
+ test("INSERT INTO ... SELECT is visible to subsequent reads") {
+ assume(isSpark40Plus, "Requires the WriteFilesExec seam")
+ // https://github.com/apache/datafusion-comet/issues/3521 - reads returned
no rows because the
+ // bespoke write path never refreshed the catalog cache. Spark's command
does that itself.
+ withTable("comet_write_target", "comet_write_source") {
+ withNativeWriter {
+ sql("CREATE TABLE comet_write_source(id bigint, name string) USING
parquet")
+ sql("CREATE TABLE comet_write_target(id bigint, name string) USING
parquet")
+ }
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ sql("INSERT INTO comet_write_source VALUES (1, 'a'), (2, 'b')")
+ }
+ withNativeWriter {
+ // 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<counter> suffix. Comet asks the commit protocol for one file per
task, so it must decline
+ // rather than quietly produce a different file layout.
+ Seq(
+ "spark.sql.files.maxRecordsPerFile" -> ((df: DataFrame, p: String) =>
df.write.parquet(p)),
+ // The write option takes precedence over the conf in FileFormatWriter,
so it must be
+ // honored here too - with the conf left at its default of 0.
+ "maxRecordsPerFile-option" -> ((df: DataFrame, p: String) =>
+ df.write.option("maxRecordsPerFile", "10").parquet(p))).foreach { case
(label, write) =>
+ withTempPath { dir =>
+ val outputPath = new File(dir, "output.parquet").getAbsolutePath
+ val sourcePath = new File(dir, "source.parquet").getAbsolutePath
+ withNativeWriter {
+ val df = materializeAsCometSource(
+ (1 to 100).map(i => (i, s"str_$i")).toDF("id",
"name").repartition(1),
+ sourcePath)
+ val confs =
+ if (label == "spark.sql.files.maxRecordsPerFile") Seq(label ->
"10") else Seq.empty
+ withSQLConf(confs: _*) {
+ val plan = captureWritePlan(p => write(df, p), outputPath)
+ assertNoCometNativeWriteExec(plan)
+ }
+ checkAnswer(spark.read.parquet(outputPath), df)
+ }
+ // Spark's writer rolled the 100 rows of the single partition into 10
files of 10 rows.
+ assert(
+ listPartFileNames(outputPath).size == 10,
+ s"$label: expected 10 rolled part files, got
${listPartFileNames(outputPath)}")
+ }
+ }
+ }
+
+ test("empty input still writes a schema-only file (SPARK-23271)") {
+ assume(isSpark40Plus, "Requires the WriteFilesExec seam")
+ // An empty input must still leave a schema behind for downstream readers:
`spark.read.parquet`
+ // of the output must see the write's schema, not fail. Comet reaches this
in two ways - if the
+ // native child has one partition producing no batches, the partition-0
branch of executeTask
+ // writes a metadata-only file; if it produces zero partitions,
doExecuteWrite swaps in a dummy
+ // single-partition RDD to get to the same branch. This test exercises the
reachable path
+ // (filtered Comet scan yielding an empty batch iterator); the
zero-partition swap is defensive
+ // because CometWriteFiles.requiresNativeChildren rules out the sources
(LocalTableScan) that
+ // would otherwise produce a zero-partition RDD.
+ withTempPath { dir =>
+ val outputPath = new File(dir, "output.parquet").getAbsolutePath
+ val sourcePath = new File(dir, "source.parquet").getAbsolutePath
+ withNativeWriter {
+ val empty = materializeAsCometSource(
+ (1 to 100).map(i => (i, s"str_$i")).toDF("id", "name"),
+ sourcePath).where("id < 0")
+ val plan = captureWritePlan(p => empty.write.parquet(p), outputPath)
+ assertHasCometNativeWriteExec(plan)
+
+ val partFiles = listPartFileNames(outputPath)
+ assert(partFiles.size == 1, s"Expected exactly one schema-only part
file, got $partFiles")
+ // Reading without an explicit schema is the point: the file must
carry it.
+ val readBack = spark.read.parquet(outputPath)
+ assert(readBack.count() == 0L)
+ assert(readBack.schema.map(_.name) == Seq("id", "name"))
+ }
+ }
+ }
+
+ test("a failing task aborts, cleans up its staging file, and the retry
succeeds") {
Review Comment:
this test actually performs a second whole write after an injected commit
failure; it does not demonstrate speculative attempts or an automatic task
retry.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
