andygrove commented on code in PR #6940:
URL: https://github.com/apache/datafusion-comet/pull/6940#discussion_r4392548798
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala:
##########
@@ -280,6 +299,29 @@ object CometWriteFilesExec extends Logging {
})
}
+ /**
+ * Runs `f`, failing the task the way Spark's writer does when `f` fails, as
+ * `FileFormatDataWriter.enrichWriteError` does on Spark 4.0+ and
`FileFormatWriter.executeTask`
+ * on 3.x: a `FetchFailedException` as it was, so that Spark reruns the
stage that produced the
+ * lost shuffle output, a `FileAlreadyExistsException` as a
+ * `TaskOutputFileAlreadyExistException` when
`spark.sql.execution.fastFailOnFileFormatOutput`
+ * is set, which fails the job without retrying the task, and anything else
as
+ * `TASK_WRITE_FAILED` for `path`.
+ *
+ * Spark's writer creates a task's file before this applies. So does Comet
for a file it creates
+ * through Hadoop, but the native writer creates local and HDFS files during
the write, so a
+ * failure to create one of those is enriched here and not in Spark.
+ */
+ private[comet] def enrichWriteError[T](path: => String)(f: => T): T =
+ try {
+ f
+ } catch {
+ case e: FetchFailedException => throw e
+ case e: FileAlreadyExistsException if
SQLConf.get.fastFailFileFormatOutput =>
Review Comment:
You are right that the branch had no test, but a pre-created target does not
reach it. The create happens in `HadoopFileOutput.create`, outside
`enrichWriteError`, and that matches Spark: its `SingleDirectoryDataWriter`
constructor creates the file before the handler
(`FileFormatDataWriter.enrichWriteError` on 4.1, the `try` in
`FileFormatWriter.executeTask` on 3.5). The `FileAlreadyExistsException` is
thrown as it is whatever `fastFailOnFileFormatOutput` says. `a task file that
already exists fails the task as Spark's writer fails it` pins that, with the
conf on and off.
The branch is reached when the exception comes from the write, the close or
the task commit. `a file that exists at task commit fails the task, without a
retry if so configured` makes `commitTask` throw it, and checks that the job
fails with `TaskOutputFileAlreadyExistException`'s message when the conf is
set, and with `TASK_WRITE_FAILED` when it is not. (5dd2fb9036)
##########
spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala:
##########
@@ -43,7 +45,58 @@ abstract class CometParquetWriterTestBase extends
CometTestBase {
SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax")(f)
}
+ // Persist `df` to `sourcePath` with Comet disabled and return a DataFrame
that reads it back,
+ // so the source plan is a Comet scan (satisfying
CometExecRule.requiresNativeChildren).
+ protected def materializeAsCometSource(df: DataFrame, sourcePath: String):
DataFrame = {
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ df.write.parquet(sourcePath)
+ }
+ spark.read.parquet(sourcePath)
+ }
+
/** The opt-in config key for native writes, which moved with the operator
on Spark 4.0+. */
+ /**
+ * Runs `write` with a source of `rows` rows in `partitions` partitions, see
the overload below.
+ */
+ protected def withRecordingFileSystem(rows: Int, partitions: Int)(
+ write: (DataFrame, String) => Unit): Unit = {
+ val session = spark
+ import session.implicits._
+ withRecordingFileSystem(
+ (1 to rows).map(i => (i, s"str_$i")).toDF("id",
"name").repartition(partitions))(write)
+ }
+
+ /**
+ * Runs `write` with `source` read natively from a local copy, 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.
+ *
+ * The output directory's name has a space, `%`, `#`, `?` and a non-ASCII
character, which a
+ * path has to keep as it passes from the commit protocol to Hadoop.
+ */
+ protected def withRecordingFileSystem(source: => DataFrame)(
Review Comment:
Added `nestedAndTemporalSource`: a struct holding an array and a map, a
nullable struct and nullable strings, decimals, a date, and a timestamp and a
timestamp_ntz, under the harness's `America/Halifax` session time zone. `nested
and temporal columns are written through Hadoop as Spark's writer writes them`
writes about 6 MiB of it in several chunks. `a write of nested and temporal
columns that fails partway is aborted` fails the third write of the same source
and checks the abort. `the default codec and several tasks match Spark's writer
through Hadoop` covers snappy and three tasks. (5dd2fb9036)
##########
spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala:
##########
@@ -1997,6 +2003,259 @@ class CometParquetWriterSuite extends
CometParquetWriterTestBase {
}
}
+ 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:
Both added (5dd2fb9036).
1. `assertMatchesSparkWriter` writes the same DataFrame with Spark's own
writer (Comet off) and checks that the native output has the same schema and
rows. The three new nested and temporal tests use it. I did not compare bytes:
the footers differ in `created_by`, so a byte comparison would fail on every
file.
2. `DataFrameWriter` never plans a V2 file write (SPARK-28396), and
`df.writeTo` into a file table fails with "Cannot write into v1 table". The one
V2 path is an `INSERT` into a table of the V2 file source, which Spark's
`FallBackFileSourceV2` plans as a V1 write. `an INSERT into a table of the V2
file source is written natively through Hadoop` covers it, with
`spark.sql.sources.useV1SourceList` empty.
##########
native/jni-bridge/src/output_stream_sink.rs:
##########
@@ -0,0 +1,141 @@
+// 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.
+
+use std::fmt::{self, Debug, Formatter};
+
+use crate::{check_exception, errors::CometError, JVMClasses};
+use datafusion::common::{DataFusionError, Result};
+use jni::objects::{Global, JByteArray, JMethodID, JObject, JValue};
+use jni::signature::{Primitive, ReturnType};
+use jni::Env;
+
+/// Receives the bytes of a file that the JVM created, in order.
+///
+/// The JVM creates the file, closes it once the native writer has written all
of it, and aborts
+/// it when the write fails, so an implementation only ever sees the file's
bytes. One writer
+/// writes one file, so calls never overlap.
+pub trait OutputStreamSink: Send + Sync {
+ /// The most bytes one call to [`Self::write`] may carry.
+ fn max_write_size(&self) -> usize;
+
+ /// Appends `data`, at most [`Self::max_write_size`] bytes, to the file.
+ fn write(&self, data: &[u8]) -> Result<()>;
+}
+
+impl Debug for dyn OutputStreamSink {
+ fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+ f.debug_struct("OutputStreamSink")
+ .field("max_write_size", &self.max_write_size())
+ .finish()
+ }
+}
+
+/// Writes to a JVM `org.apache.comet.parquet.HadoopFileOutput` from any
native thread.
+///
+/// Each write copies its bytes into a Java array this sink allocates once,
then hands the array
+/// to `write(byte[], int)`, so a write costs one copy and one upcall but no
Java allocation. The
+/// global references keep the Java object, its class and the array alive for
as long as the
+/// cached method ID is used.
+pub struct JavaOutputStreamSink {
+ output: Global<JObject<'static>>,
+ buffer: Global<JByteArray<'static>>,
+ write_method: JMethodID,
+}
+
+impl JavaOutputStreamSink {
+ /// Bytes per upcall, large enough that the upcall costs little next to
the copy it carries.
+ pub const MAX_WRITE_SIZE: usize = 1024 * 1024;
+
+ /// Captures `output` while running on an attached JVM thread.
+ pub fn try_new(env: &mut Env<'_>, output: &JObject<'_>) -> Result<Self> {
+ if output.is_null() {
+ return Err(DataFusionError::Execution(
+ "The output stream of a native write must not be
null".to_string(),
+ ));
+ }
+
+ let output_class =
env.get_object_class(output).map_err(CometError::from)?;
+ let write_method = env
+ .get_method_id(
+ &output_class,
+ jni::jni_str!("write"),
+ jni::jni_sig!("([BI)V"),
+ )
+ .map_err(CometError::from)?;
+ let buffer = env
+ .new_byte_array(Self::MAX_WRITE_SIZE)
+ .map_err(CometError::from)?;
+
+ Ok(Self {
+ output: env.new_global_ref(output).map_err(CometError::from)?,
+ buffer: env.new_global_ref(&buffer).map_err(CometError::from)?,
+ write_method,
+ })
+ }
+}
+
+impl OutputStreamSink for JavaOutputStreamSink {
+ fn max_write_size(&self) -> usize {
+ Self::MAX_WRITE_SIZE
+ }
+
+ fn write(&self, data: &[u8]) -> Result<()> {
+ if data.len() > Self::MAX_WRITE_SIZE {
Review Comment:
Added three Rust tests (5dd2fb9036). The two checks needed a JVM, so I moved
them into `check_not_null` and `checked_write_length`, as
`checked_payload_length` is for the shuffle pusher. The tests cover a null
output being rejected, an oversize write failing with `Internal` for sizes up
to `usize::MAX`, and the limit itself being accepted.
--
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]
