andygrove commented on code in PR #6940:
URL: https://github.com/apache/datafusion-comet/pull/6940#discussion_r4342758731
##########
spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala:
##########
@@ -187,6 +189,117 @@ object NativeWriteUtils {
"spark.comet.write.parquet.enabled=false to write this table with
Spark.")
}
+ /**
+ * Whether the native writer creates the file at `filePath` itself, which it
does for local and
+ * HDFS files. Any other file is created through Hadoop's FileSystem API and
the native writer
+ * only writes its bytes, see [[org.apache.comet.parquet.HadoopFileOutput]].
That is always
+ * correct, whatever the file system and whatever the committer expects of
it, which is why this
+ * is decided per task on the path the commit protocol actually returned
rather than at
+ * planning.
+ */
+ def nativeWriterCreatesFile(filePath: String): Boolean =
+ filePath.startsWith("file:") || filePath.startsWith("hdfs:")
+
+ /**
+ * A fallback reason when the native writer cannot write to `outputPath`'s
file system, or
+ * `None` when it can. Local and HDFS files are created natively, and so are
the files of a
Review Comment:
Fixed, thanks. It now says local and HDFS files are created natively, and
the files the magic committer commits are created through S3A and written
natively. (50efcb6ec)
##########
native/core/src/execution/jni_api.rs:
##########
@@ -929,23 +938,80 @@ fn register_shuffle_partition_pusher(
exec_context: &mut ExecutionContext,
make_pusher: impl FnOnce() -> CometResult<Arc<dyn ShufflePartitionPusher>>,
) -> CometResult<()> {
- if exec_context.root_op.is_some() {
- return Err(CometError::Internal(
- "Remote shuffle callback cannot be registered after native
execution starts"
- .to_string(),
- ));
+ register_before_execution(
+ exec_context.root_op.is_some(),
+ &mut exec_context.shuffle_partition_pusher,
+ "Remote shuffle callback",
+ make_pusher,
+ )
+}
+
+/// Registers one of a plan's task-owned JVM callbacks in `slot`, built by
`make` only once the
+/// plan is known to accept it: before execution has `started`, and once.
`what` names the callback
+/// in the errors.
+fn register_before_execution<T: ?Sized>(
+ started: bool,
+ slot: &mut Option<Arc<T>>,
+ what: &str,
+ make: impl FnOnce() -> CometResult<Arc<T>>,
+) -> CometResult<()> {
+ if started {
+ return Err(CometError::Internal(format!(
+ "{what} cannot be registered after native execution starts"
+ )));
}
- if exec_context.shuffle_partition_pusher.is_some() {
- return Err(CometError::Internal(
- "Remote shuffle callback has already been registered for this
task".to_string(),
- ));
+ if slot.is_some() {
+ return Err(CometError::Internal(format!(
+ "{what} has already been registered for this task"
+ )));
}
- exec_context.shuffle_partition_pusher = Some(make_pusher()?);
+ *slot = Some(make()?);
Ok(())
}
+/// Binds the output stream of the file the plan's Parquet writer writes, for
a file the JVM
+/// created itself, before native execution is initialized.
+///
+/// Separate from `createPlan` for the same reason as
`setShufflePartitionPusher`: only native
+/// writes the JVM creates the file for need it, so every other caller keeps
the existing ABI.
+#[no_mangle]
+pub extern "system" fn Java_org_apache_comet_Native_setOutputStream(
+ e: EnvUnowned,
+ _class: JClass,
+ exec_context: jlong,
+ output: JObject,
+) {
+ try_unwrap_or_throw(&e, |env| {
+ if exec_context == 0 {
+ return Err(CometError::NullPointer(
+ "A native write's output stream requires a valid native
execution plan".to_string(),
+ ));
+ }
+
+ let exec_context = get_execution_context(exec_context);
+ register_output_stream(exec_context, || {
+ Ok(Arc::new(JavaOutputStreamSink::try_new(env, &output)?))
+ })
+ })
+}
+
+/// Registers the output stream of the file the plan's Parquet writer writes,
built by
+/// `make_sink` only once the plan is known to accept one: before execution
starts, and once. Core
+/// of `Native.setOutputStream`, whose sink calls back into the JVM.
+fn register_output_stream(
+ exec_context: &mut ExecutionContext,
+ make_sink: impl FnOnce() -> CometResult<Arc<dyn OutputStreamSink>>,
+) -> CometResult<()> {
+ register_before_execution(
+ exec_context.root_op.is_some(),
+ &mut exec_context.output_stream,
+ "A native write's output stream",
+ make_sink,
+ )
+}
Review Comment:
Done. `setOutputStream` now calls `register_before_execution` directly, and
the module doc no longer lists `register_output_stream`. The once and
before-execution rules stay covered by the pusher test, which goes through the
same helper. (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,
Review Comment:
Done. `writeFile` now takes the template `Operator` and the caller's
`ParquetWriter.Builder`. It sets `output_path` and `jvm_output_stream` on the
builder and builds the task's plan once.
It also creates the iterator through `CometExec.getCometIterator` again. The
`Operator` overload now takes an `outputStream`, but without a default. The
serialized-plan overload already has one (`outputTypes`), and Scala allows
default arguments on only one overload of a method. So its three other callers
pass `None`. (50efcb6ec)
##########
docs/source/user-guide/latest/operators.md:
##########
@@ -140,6 +140,14 @@ combined with `OR`, which Spark plans as an existence
`BroadcastNestedLoopJoinEx
| Iceberg writes: `AppendDataExec`, `OverwriteByExpressionExec`,
`OverwritePartitionsDynamicExec`, `ReplaceDataExec` | ⚠️ | Disabled by
default (opt-in). `spark.comet.write.iceberg.enabled=true` plans the write as
`IcebergWrite` and `IcebergCommit` and writes eligible data files natively
(`CometIcebergWrite`). Other writes use Iceberg's Java writer. Covers `INSERT
INTO`, `INSERT OVERWRITE`, and copy-on-write `DELETE` / `UPDATE` / `MERGE`.
Merge-on-read writes (`WriteDeltaExec`) fall back. See [Iceberg
Writes](iceberg.md#writes). |
| `MergeRowsExec`
| ⚠️ | Spark 3.5+. Experimental,
disabled by default (`spark.comet.exec.mergeRows.enabled`). On Spark 4.1+,
stock V2 writers retain Spark MergeRowsExec; Comet's split Iceberg path can run
it natively while preserving MergeSummary. See [MERGE
INTO](compatibility/operators.md#merge-into-mergerowsexec).
|
+Native Parquet writes support local and HDFS destinations, and S3 writes that
the S3A magic
+committer commits. Spark 4.1 configures that committer by default when
spark-hadoop-cloud is on the
+classpath, and earlier versions for a bucket that sets
+`fs.s3a.bucket.<bucket>.committer.magic.enabled`. S3A creates each of those
files, through the same
+call Spark's own writer makes, so the magic committer and S3A's write
settings, such as server-side
+encryption, ACLs, storage class and object headers, apply exactly as they do
for Spark's writer.
+Other S3 writes fall back to Spark.
Review Comment:
Moved. The paragraph is now in the Writing section of `parquet.md`, right
after the scope sentence, and both write rows in `operators.md` link to
`parquet.md#writing`. (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]
