parthchandra commented on code in PR #6940:
URL: https://github.com/apache/datafusion-comet/pull/6940#discussion_r4392274964
##########
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:
This fast-fail branch turns a `FileAlreadyExistsException` into
`TaskOutputFileAlreadyExistException` when
`spark.sql.execution.fastFailOnFileFormatOutput` is set, but nothing tests it.
`RecordingFileSystem` already creates files with `overwrite=false`, so a test
that pre-creates the target path and sets that conf would cover this branch.
Right now `TaskOutputFileAlreadyExistException` is only imported, never
asserted.
##########
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:
Two things on these RecordingFileSystem tests:
1. They check the file reads back equal to the input DataFrame, which is a
round-trip check, not a comparison against Spark's own writer. For at least one
case it would be good to assert the output matches Spark (either
`checkSparkAnswer`, or byte-identical to Spark's Parquet writer), so we catch a
case where Comet and Spark write the same data differently.
2. They all go through `df.write.parquet`, which is the V1 file-source path.
The DataSource V2 write path isn't exercised through the Hadoop output-stream
route. Worth one V2 write so both planning paths are covered.
##########
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:
Every write through this new Hadoop output-stream path uses simple `(id,
name)` or random-string data. The complex-type and temporal coverage (structs
with map and array fields, nullable and non-nullable, decimals, dates,
timestamps, timestamp_ntz, timezones) lives in the local-file tests, so it
never runs through this path. Could we add at least one complex-typed source
and one timestamp / timestamp_ntz source here? The path just forwards Arrow
bytes, so the risk is low, but nothing proves a nested or temporal file
survives the chunked 1 MiB writes and the abort path.
##########
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:
This oversize guard (`data.len() > MAX_WRITE_SIZE`) and the null-output
check near line 65 are never hit by a test. This file has no unit tests, and
the Scala integration tests only drive the normal path. A small Rust test for
each (an oversize write returns the `Internal` error, a null output is
rejected) would lock in these defensive branches.
--
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]
