andygrove commented on code in PR #6940:
URL: https://github.com/apache/datafusion-comet/pull/6940#discussion_r4342759616
##########
spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala:
##########
@@ -1997,6 +2002,175 @@ class CometParquetWriterSuite extends
CometParquetWriterTestBase {
}
}
+ /**
+ * Runs `write` with a source of `rows` rows in `partitions` partitions,
read natively, and the
+ * path to write it to, while RecordingFileSystemCommitProtocol puts the
write's task files on
+ * [[RecordingFileSystem]], whose recordings start empty and are cleared
afterwards.
+ */
+ private def withRecordingFileSystem(rows: Int, partitions: Int)(
+ write: (DataFrame, String) => Unit): Unit =
+ withTempPath { dir =>
+ val outputPath = new File(dir, "output.parquet").getAbsolutePath
+ val sourcePath = new File(dir, "source.parquet").getAbsolutePath
+ withNativeWriter {
+ val df = materializeAsCometSource(
+ (1 to rows).map(i => (i, s"str_$i")).toDF("id",
"name").repartition(partitions),
+ sourcePath)
+ RecordingFileSystem.reset()
+ try {
+ withSQLConf(
+ SQLConf.FILE_COMMIT_PROTOCOL_CLASS.key ->
+ classOf[RecordingFileSystemCommitProtocol].getName,
+ s"fs.${RecordingFileSystem.SCHEME}.impl" ->
classOf[RecordingFileSystem].getName) {
+ write(df, outputPath)
+ }
+ } finally {
+ RecordingFileSystem.reset()
+ }
+ }
+ }
+
+ test("a task file the native writer cannot create is created through
Hadoop") {
+ // The native writer creates only local and HDFS files. Every other file
is created through
+ // Hadoop's FileSystem API (see HadoopFileOutput) and the native writer
only supplies the
+ // bytes. RecordingFileSystemCommitProtocol puts task files on a scheme
the native writer
+ // cannot open at all, so the write can only succeed that way.
+ withRecordingFileSystem(rows = 1000, partitions = 3) { (df, outputPath) =>
Review Comment:
Added `a file of several MiB is written through Hadoop in writes of at most
1 MiB`. It writes 30,000 rows of random two-byte UTF-8 strings with
`compression=none`, about 6 MiB in one file. It checks that the data reads
back, that every write is at most 1 MiB, that there are at least as many writes
as MiB, and that the writes add up to the file's size.
`RecordingFileSystem.failAtWrite` now fails the Nth write. `a write that
fails partway through a file aborts what it already wrote` fails the third one,
after two MiB have gone through, and checks that the file is aborted and
nothing is left. (50efcb6ec)
##########
spark/src/test/java/org/apache/comet/hadoop/fs/RecordingFileSystem.java:
##########
@@ -0,0 +1,175 @@
+/*
+ * 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.hadoop.fs;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.net.URI;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+
+import org.apache.hadoop.fs.Abortable;
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.RawLocalFileSystem;
+import org.apache.hadoop.fs.StreamCapabilities;
+import org.apache.hadoop.util.Progressable;
+
+/**
+ * A local-disk-backed FileSystem under its own scheme that records every file
created through it
+ * and how each file's stream ended, and can fail writes on demand. Its
streams can abort, as S3A's
+ * do, and an aborted file is deleted, as an aborted S3A upload never becomes
an object.
+ *
+ * <p>It lets a test check that a file was created through Hadoop's FileSystem
API without a remote
+ * store: Comet's native writer creates only local and HDFS files itself. The
recordings are static,
+ * since FileSystem instances are cached and the tests run in local mode.
+ */
+public class RecordingFileSystem extends RawLocalFileSystem {
+
+ public static final String SCHEME = "comet-recording";
+
+ public static final String PREFIX = SCHEME + "://recording";
+
+ /** The message of the {@link IOException} every write throws while {@link
#failWrites} is set. */
+ public static final String WRITE_FAILURE = "injected output stream failure";
+
+ /** The files created through this FileSystem, in order. */
+ public static final List<String> created = new CopyOnWriteArrayList<>();
+
+ /** The files whose stream was closed. */
+ public static final List<String> closed = new CopyOnWriteArrayList<>();
+
+ /** The files whose stream was aborted. */
+ public static final List<String> aborted = new CopyOnWriteArrayList<>();
+
+ public static volatile boolean failWrites = false;
+
+ public static void reset() {
+ created.clear();
+ closed.clear();
+ aborted.clear();
+ failWrites = false;
+ }
+
+ public RecordingFileSystem() {
+ // Avoid `URI scheme is not "file"` error on
+ // RawLocalFileSystem$DeprecatedRawLocalFileStatus.getOwner
+ RawLocalFileSystem.useStatIfAvailable();
+ }
+
+ @Override
+ public String getScheme() {
+ return SCHEME;
+ }
+
+ @Override
+ public URI getUri() {
+ return URI.create(PREFIX);
+ }
+
+ @Override
+ public FSDataOutputStream create(
+ Path f,
+ boolean overwrite,
+ int bufferSize,
+ short replication,
+ long blockSize,
+ Progressable progress)
+ throws IOException {
+ return record(f, super.create(f, overwrite, bufferSize, replication,
blockSize, progress));
+ }
+
+ private FSDataOutputStream record(Path path, FSDataOutputStream out) {
+ created.add(path.toString());
+ return new FSDataOutputStream(new RecordingStream(path, out), statistics);
+ }
+
+ private final class RecordingStream extends OutputStream
+ implements Abortable, StreamCapabilities {
+ private final Path path;
+ private final OutputStream out;
+ private boolean finished = false;
+
+ RecordingStream(Path path, OutputStream out) {
+ this.path = path;
+ this.out = out;
+ }
+
+ @Override
+ public void write(int b) throws IOException {
+ checkWritable();
+ out.write(b);
+ }
+
+ @Override
+ public void write(byte[] b, int off, int len) throws IOException {
+ checkWritable();
+ out.write(b, off, len);
+ }
+
+ private void checkWritable() throws IOException {
+ if (failWrites) {
+ throw new IOException(WRITE_FAILURE);
+ }
+ }
+
+ @Override
+ public void close() throws IOException {
Review Comment:
Added `failCloses` and `a file whose close fails fails its task, which is
aborted`. The task fails with `TASK_WRITE_FAILED` caused by the close's
`IOException`. `abortTask` runs once, and no part file or `_temporary` is left.
The file is not aborted after a failed close. That matches Spark's writer,
whose `releaseCurrentWriter` closes the writer once and clears it, so the abort
that follows never touches the file again. The commit protocol's `abortTask`
removes what is left. (50efcb6ec)
##########
spark/src/test/java/org/apache/comet/hadoop/fs/RecordingFileSystem.java:
##########
@@ -0,0 +1,175 @@
+/*
+ * 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.hadoop.fs;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.net.URI;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+
+import org.apache.hadoop.fs.Abortable;
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.RawLocalFileSystem;
+import org.apache.hadoop.fs.StreamCapabilities;
+import org.apache.hadoop.util.Progressable;
+
+/**
+ * A local-disk-backed FileSystem under its own scheme that records every file
created through it
+ * and how each file's stream ended, and can fail writes on demand. Its
streams can abort, as S3A's
+ * do, and an aborted file is deleted, as an aborted S3A upload never becomes
an object.
+ *
+ * <p>It lets a test check that a file was created through Hadoop's FileSystem
API without a remote
+ * store: Comet's native writer creates only local and HDFS files itself. The
recordings are static,
+ * since FileSystem instances are cached and the tests run in local mode.
+ */
+public class RecordingFileSystem extends RawLocalFileSystem {
+
+ public static final String SCHEME = "comet-recording";
+
+ public static final String PREFIX = SCHEME + "://recording";
+
+ /** The message of the {@link IOException} every write throws while {@link
#failWrites} is set. */
+ public static final String WRITE_FAILURE = "injected output stream failure";
+
+ /** The files created through this FileSystem, in order. */
+ public static final List<String> created = new CopyOnWriteArrayList<>();
+
+ /** The files whose stream was closed. */
+ public static final List<String> closed = new CopyOnWriteArrayList<>();
+
+ /** The files whose stream was aborted. */
+ public static final List<String> aborted = new CopyOnWriteArrayList<>();
+
+ public static volatile boolean failWrites = false;
+
+ public static void reset() {
+ created.clear();
+ closed.clear();
+ aborted.clear();
+ failWrites = false;
+ }
+
+ public RecordingFileSystem() {
+ // Avoid `URI scheme is not "file"` error on
+ // RawLocalFileSystem$DeprecatedRawLocalFileStatus.getOwner
+ RawLocalFileSystem.useStatIfAvailable();
+ }
+
+ @Override
+ public String getScheme() {
+ return SCHEME;
+ }
+
+ @Override
+ public URI getUri() {
+ return URI.create(PREFIX);
+ }
+
+ @Override
+ public FSDataOutputStream create(
+ Path f,
+ boolean overwrite,
+ int bufferSize,
+ short replication,
+ long blockSize,
+ Progressable progress)
+ throws IOException {
+ return record(f, super.create(f, overwrite, bufferSize, replication,
blockSize, progress));
+ }
+
+ private FSDataOutputStream record(Path path, FSDataOutputStream out) {
+ created.add(path.toString());
+ return new FSDataOutputStream(new RecordingStream(path, out), statistics);
+ }
+
+ private final class RecordingStream extends OutputStream
+ implements Abortable, StreamCapabilities {
+ private final Path path;
+ private final OutputStream out;
+ private boolean finished = false;
+
+ RecordingStream(Path path, OutputStream out) {
+ this.path = path;
+ this.out = out;
+ }
+
+ @Override
+ public void write(int b) throws IOException {
+ checkWritable();
+ out.write(b);
+ }
+
+ @Override
+ public void write(byte[] b, int off, int len) throws IOException {
+ checkWritable();
+ out.write(b, off, len);
+ }
+
+ private void checkWritable() throws IOException {
+ if (failWrites) {
+ throw new IOException(WRITE_FAILURE);
+ }
+ }
+
+ @Override
+ public void close() throws IOException {
Review Comment:
Added an `abortable` switch. `a failed write closes a file whose stream
cannot abort` checks that the file is closed rather than aborted, that
`abortTask` runs, and that nothing is left. (50efcb6ec)
##########
spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala:
##########
@@ -1997,6 +2002,175 @@ class CometParquetWriterSuite extends
CometParquetWriterTestBase {
}
}
+ /**
+ * Runs `write` with a source of `rows` rows in `partitions` partitions,
read natively, and the
+ * path to write it to, while RecordingFileSystemCommitProtocol puts the
write's task files on
+ * [[RecordingFileSystem]], whose recordings start empty and are cleared
afterwards.
+ */
+ private def withRecordingFileSystem(rows: Int, partitions: Int)(
+ write: (DataFrame, String) => Unit): Unit =
+ withTempPath { dir =>
+ val outputPath = new File(dir, "output.parquet").getAbsolutePath
+ val sourcePath = new File(dir, "source.parquet").getAbsolutePath
+ withNativeWriter {
+ val df = materializeAsCometSource(
+ (1 to rows).map(i => (i, s"str_$i")).toDF("id",
"name").repartition(partitions),
+ sourcePath)
+ RecordingFileSystem.reset()
+ try {
+ withSQLConf(
+ SQLConf.FILE_COMMIT_PROTOCOL_CLASS.key ->
+ classOf[RecordingFileSystemCommitProtocol].getName,
+ s"fs.${RecordingFileSystem.SCHEME}.impl" ->
classOf[RecordingFileSystem].getName) {
+ write(df, outputPath)
+ }
+ } finally {
+ RecordingFileSystem.reset()
+ }
+ }
+ }
+
+ test("a task file the native writer cannot create is created through
Hadoop") {
+ // The native writer creates only local and HDFS files. Every other file
is created through
+ // Hadoop's FileSystem API (see HadoopFileOutput) and the native writer
only supplies the
+ // bytes. RecordingFileSystemCommitProtocol puts task files on a scheme
the native writer
+ // cannot open at all, so the write can only succeed that way.
+ withRecordingFileSystem(rows = 1000, partitions = 3) { (df, outputPath) =>
+ val plan = captureWritePlan(p => df.write.parquet(p), outputPath)
+ assertHasCometNativeWriteExec(plan)
+
+ val created = RecordingFileSystem.created.asScala.toSeq
+ assert(created.size == 3, s"Expected one file per task, got $created")
+ assert(created.forall(_.startsWith(RecordingFileSystem.PREFIX)), created)
+ // Closed, which is where a FileSystem such as S3A completes or records
the upload, and
+ // only once the native writer had written all of it.
+ assert(RecordingFileSystem.closed.asScala.toSet == created.toSet)
+ assert(RecordingFileSystem.aborted.isEmpty)
+
+ checkAnswer(spark.read.parquet(outputPath), df)
+ assert(listPartFileNames(outputPath).size == 3)
+ assert(new File(outputPath, "_SUCCESS").exists())
+ }
+ }
+
+ test("a failed write through Hadoop aborts the file and fails with what its
stream threw") {
Review Comment:
Added `CometParquetWriterRetrySuite`. It runs under `local[5,2]`, so a task
gets a second attempt, and `RecordingFileSystem.failingFiles` limits the
injected failure to the first attempt's file. It checks one aborted file, one
closed file and one `abortTask`, each row once in the output, one part file,
and no `_temporary`.
It is a suite of its own because the other failure tests expect a failed
task to fail the job. It is registered in both PR workflows. (50efcb6ec)
##########
spark/src/test/scala/org/apache/comet/parquet/ParquetWriteWithMagicCommitterSuite.scala:
##########
@@ -0,0 +1,228 @@
+/*
+ * 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.parquet
+
+import java.io.File
+
+import scala.jdk.CollectionConverters._
+import scala.util.Try
+
+import org.apache.hadoop.fs.Path
+import org.apache.hadoop.fs.s3a.commit.files.SuccessData
+import org.apache.parquet.hadoop.ParquetOutputCommitter
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.comet.CometNativeWriteExec
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.command.DataWritingCommandExec
+import org.apache.spark.sql.functions.col
+import org.apache.spark.sql.internal.SQLConf
+
+import org.apache.comet.CometS3TestBase
+import org.apache.comet.serde.operator.NativeWriteUtils
+
+import software.amazon.awssdk.services.s3.model.ListMultipartUploadsRequest
+
+/**
+ * Native Parquet writes to S3 that S3A's magic committer commits, against
MinIO.
+ *
+ * A manual suite like the other MinIO suites, since it needs Docker. Every
session here gets the
+ * magic committer the way Spark configures it when spark-hadoop-cloud is on
the classpath: by
+ * default on Spark 4.1+ (SPARK-47618), and for a bucket that asks for it on
earlier versions. The
+ * shared test configuration turns that off for every other suite.
+ */
+class ParquetWriteWithMagicCommitterSuite
+ extends CometParquetWriterTestBase
+ with CometS3TestBase {
+
+ import testImplicits._
+
+ override protected val testBucketName = "magic-committer"
+
+ override protected def sparkConf: SparkConf = {
+ val conf = super.sparkConf
+ conf.set("spark.hadoop.fs.s3a.committer.magic.enabled", "true")
+
conf.set(s"spark.hadoop.fs.s3a.bucket.$testBucketName.committer.magic.enabled",
"true")
+ conf
+ }
+
+ /** The multipart uploads still pending in `bucket`, as their object keys. */
+ private def pendingUploads(bucket: String): Seq[String] = withS3Client {
client =>
+ client
+
.listMultipartUploads(ListMultipartUploadsRequest.builder().bucket(bucket).build())
+ .uploads()
+ .asScala
+ .map(_.key())
+ .toSeq
+ }
+
+ /** `rows` rows read natively from a local copy, so the write's input is a
Comet scan. */
+ private def source(rows: Int, dir: File): DataFrame =
+ materializeAsCometSource(
+ (1 to rows).map(i => (i, s"str_$i")).toDF("id", "name"),
+ new File(dir, "source.parquet").getAbsolutePath)
+
+ /**
+ * The files and bytes the write reports: from Spark's stats tracker on
Spark 4.0+, which leaves
+ * the write command in the plan, and from the native writer on 3.x, which
replaces it.
+ */
+ private def writtenFilesAndBytes(plan: SparkPlan): (Long, Long) =
+ plan
+ .collectFirst {
+ case write: DataWritingCommandExec =>
+ (write.cmd.metrics("numFiles").value,
write.cmd.metrics("numOutputBytes").value)
+ case write: CometNativeWriteExec =>
+ (write.metrics("files_written").value,
write.metrics("bytes_written").value)
+ }
+ .getOrElse(fail(s"No write in the plan:\n${plan.treeString}"))
+
+ /**
+ * Writes `rows` rows in `files` tasks natively to `table` under `confs`,
and checks that the
+ * magic committer committed them: its `_SUCCESS` manifest names it and the
files, no `__magic`
+ * directory or pending upload is left behind, the write reports each file's
real size, and the
+ * data reads back. On Spark 4.0+ that size comes from the marker S3A writes
when it closes a
+ * magic file, and on 3.x from what the native writer counted.
+ */
+ private def checkMagicCommittedWrite(table: String, rows: Int, files: Int)(
+ confs: (String, String)*): Unit =
+ withTempPath { dir =>
+ val df = source(rows, dir).repartition(files)
+ withNativeWriter {
+ withSQLConf(confs: _*) {
+ val plan = captureWritePlan(p => df.write.parquet(p), table)
+ assertHasCometNativeWriteExec(plan)
+
+ val tablePath = new Path(table)
+ val fs = tablePath.getFileSystem(spark.sessionState.newHadoopConf())
+ val manifest = SuccessData.load(fs, new Path(tablePath, "_SUCCESS"))
+ assert(manifest.getCommitter == "magic", manifest)
+ assert(manifest.getFilenames.size == files, manifest)
+
+ val entries = fs.listStatus(tablePath)
+ assert(
+ entries.forall(!_.getPath.getName.startsWith("__magic")),
+ entries.map(_.getPath).mkString(", "))
+ val parts = entries.filter(_.getPath.getName.startsWith("part-"))
+ assert(parts.length == files, parts.map(_.getPath).mkString(", "))
+ assert(pendingUploads(tablePath.toUri.getHost).isEmpty)
+ assert(writtenFilesAndBytes(plan) == ((files.toLong,
parts.map(_.getLen).sum)))
+ }
+ }
+ checkAnswer(spark.read.parquet(table), df)
+ }
+
+ test("a native write is committed by S3A's magic committer") {
+ // Spark's own magic committer configuration, rather than one this suite
spells out.
+ assert(
+ spark.conf.get(SQLConf.FILE_COMMIT_PROTOCOL_CLASS.key) ==
+ NativeWriteUtils.PathOutputCommitProtocolClass)
+ assert(
+ spark.conf.get(SQLConf.PARQUET_OUTPUT_COMMITTER_CLASS.key) ==
+ NativeWriteUtils.BindingParquetOutputCommitterClass)
+ checkMagicCommittedWrite(s"s3a://$testBucketName/committed", rows = 10000,
files = 4)()
+ }
+
+ test("a native write is committed with the magic committer tracking commits
in memory") {
+ assume(
+ Try(
+ Class.forName(
+
"org.apache.hadoop.fs.s3a.commit.magic.InMemoryMagicCommitTracker")).isSuccess,
+ "Requires S3A's in-memory commit tracking, which Hadoop 3.3.4 does not
have")
+ // Task commit then reads the pending uploads from S3A's in-process
tracker, not from the
+ // `.pending` files: a native writer that only reproduced those files
would lose its data
+ // here. Its own bucket, since S3A reads the setting from the file system
instance it caches
+ // per bucket.
+ val bucket = "magic-in-memory"
+ createBucketIfNotExists(bucket)
+ checkMagicCommittedWrite(s"s3a://$bucket/committed", rows = 1000, files =
2)(
+ "fs.s3a.committer.magic.track.commits.in.memory.enabled" -> "true")
+ }
+
+ // Spark 4.1 configures both bindings at once. Either one alone asks the
committer factory for
+ // the magic committer: PathOutputCommitProtocol ignores the Parquet
committer class, and Spark's
+ // default protocol instantiates BindingParquetOutputCommitter, which hands
out the magic
+ // committer's paths.
+ Seq(
+ "PathOutputCommitProtocol" -> (SQLConf.PARQUET_OUTPUT_COMMITTER_CLASS.key
->
+ classOf[ParquetOutputCommitter].getName),
+ "BindingParquetOutputCommitter" -> (SQLConf.FILE_COMMIT_PROTOCOL_CLASS.key
->
+ SQLConf.FILE_COMMIT_PROTOCOL_CLASS.defaultValueString)).foreach {
+ case (binding, otherBindingUnset) =>
+ test(s"a native write is committed by the magic committer through
$binding alone") {
+ checkMagicCommittedWrite(s"s3a://$testBucketName/$binding", rows =
1000, files = 2)(
+ otherBindingUnset)
+ }
+ }
+
+ test("a failed native write leaves no pending upload behind") {
Review Comment:
Added all three, and ran the suite against MinIO on Spark 4.1 (9 passed) and
3.5 (8 passed, with the in-memory tracking test canceled on Hadoop 3.3.4).
- `a native write that fails after S3A uploaded part of its file leaves
nothing behind` sets `fs.s3a.multipart.size` to 5M in a bucket of its own,
because S3A reads it from the file system it caches per bucket. The native
writer hands a row group to the stream once it has 1048576 rows, about 8 MiB of
uncompressed doubles here, and the division by zero comes after that. No object
or pending upload is left.
- `a native write whose task commit fails leaves nothing behind` uses a
`PathOutputCommitProtocol` subclass that throws in `commitTask` before calling
`super`. Only S3A's `abortTask` can clean up there, and it does.
- For `commitJob`, the expectation turned out not to hold for Spark's own
writer. With the same subclass throwing in `commitJob`, Spark's writer leaves
the committed task's upload pending, and so does Comet's. That happens whether
`fs.s3a.committer.abort.pending.uploads` is on or off. The magic committer's
`abortJob` removes `__magic` but does not cancel the uploads of tasks that
already committed. So `a native write whose job commit fails leaves what
Spark's writer leaves` runs the failed write with both writers. It checks that
nothing is visible and that the native writer leaves exactly as many pending
uploads as Spark's, then aborts them. (81f116636)
One change was needed to make the subclass cases run natively at all. The
planning gate matched `PathOutputCommitProtocol` by exact class name, so a
subclass fell back to Spark. It now admits subclasses too. That is safe for the
same reason as before: each task writes through Hadoop whatever path its
committer returns, whatever the committer is. The planning test covers the
subclass case.
##########
spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala:
##########
@@ -1997,6 +2002,175 @@ class CometParquetWriterSuite extends
CometParquetWriterTestBase {
}
}
+ /**
+ * Runs `write` with a source of `rows` rows in `partitions` partitions,
read natively, and the
+ * path to write it to, while RecordingFileSystemCommitProtocol puts the
write's task files on
+ * [[RecordingFileSystem]], whose recordings start empty and are cleared
afterwards.
+ */
+ private def withRecordingFileSystem(rows: Int, partitions: Int)(
+ write: (DataFrame, String) => Unit): Unit =
+ withTempPath { dir =>
+ val outputPath = new File(dir, "output.parquet").getAbsolutePath
Review Comment:
Done. The output directory in `withRecordingFileSystem` is now `output dir
%#?é.parquet`, so every test on that file system writes through it and reads it
back. All of them pass. (50efcb6ec)
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala:
##########
@@ -318,25 +327,77 @@ object CometWriteFilesExec extends Logging {
}
})
- val execIterator = CometExec.getCometIterator(
+ writeFile(
+ filePath,
+ hadoopConf,
+ taskWrite.nativeOp.toBuilder.setParquetWriter(parquetWriter).build(),
CometArrowStream.inputObjects(trackedBatches, 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()
+ partitionId)
+ }
+
+ /**
+ * Writes the task file at `filePath` by running `nativeOp`, a native write
plan whose Parquet
+ * writer is set up for the task except for the file, over `inputs`. Both
native writers write
+ * their files with this: CometWriteFilesExec on Spark 4.0+ and
CometNativeWriteExec on 3.x.
+ *
+ * When the native writer cannot create the file itself (see
+ * [[NativeWriteUtils.nativeWriterCreatesFile]]), it is created here first,
through Hadoop's
+ * FileSystem API with the task's configuration, as [[HadoopFileOutput]]
explains, and the
+ * plan's writer writes the file's bytes to it. The file is closed once the
plan has written all
+ * of it, and aborted if the plan fails.
+ */
+ private[comet] def writeFile(
+ filePath: String,
+ hadoopConf: Configuration,
+ nativeOp: Operator,
+ inputs: Array[Object],
+ numOutputCols: Int,
+ nativeMetrics: CometMetricNode,
+ numPartitions: Int,
+ partitionId: Int): Unit = {
+ def run(output: Option[HadoopFileOutput]): Unit = {
+ val writer = nativeOp.getParquetWriter.toBuilder
+ .setOutputPath(filePath)
+ .setJvmOutputStream(output.isDefined)
+ .build()
+ val execIterator = new CometExecIterator(
+ CometExec.newIterId,
+ inputs,
+ numOutputCols,
+
CometExec.serializeNativePlan(nativeOp.toBuilder.setParquetWriter(writer).build()),
+ nativeMetrics,
+ numPartitions,
+ partitionId,
+ outputStream = output)
+
+ // `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()
+ }
+ }
+
+ if (NativeWriteUtils.nativeWriterCreatesFile(filePath)) {
+ run(None)
+ } else {
+ val file = HadoopFileOutput.create(filePath, hadoopConf)
Review Comment:
Good catch. It wasn't intended, and this PR now matches Spark.
`CometWriteFilesExec.enrichWriteError` has the same three cases as
`FileFormatDataWriter.enrichWriteError`. A `FetchFailedException` is rethrown
as it was. A `FileAlreadyExistsException` becomes
`TaskOutputFileAlreadyExistException` when
`spark.sql.execution.fastFailOnFileFormatOutput` is set. Anything else becomes
`TASK_WRITE_FAILED`.
It is applied where Spark applies it. On 4.0+, a failure while writing is
reported for the file's path, and a failure while committing the task,
including closing the file for the trackers, for the write's path. On 3.x, both
are reported for the write's path, as `FileFormatWriter.executeTask` does. It
covers every native write, not just the Hadoop route.
One correction on `FileAlreadyExistsException`. Spark's writer creates the
task file in the `DataWriter` constructor, which runs before `enrichWriteError`
on 4.0+ and outside `executeTask`'s `try` on 3.x. So a create that finds an
existing file is not converted in Spark either, and the task is retried.
`HadoopFileOutput.create` is also outside the enrichment now, so the Hadoop
route matches that. The native writer creates local and HDFS files during the
write, so a failure there is enriched, and the comment on `enrichWriteError`
says so.
The failure tests now check for `TASK_WRITE_FAILED` with the stream's own
`IOException` as its cause, and the test comment no longer claims more than
that. (50efcb6ec)
--
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]
