comphead commented on code in PR #6940:
URL: https://github.com/apache/datafusion-comet/pull/6940#discussion_r4317449100
##########
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:
`register_output_stream` now has one caller, `setOutputStream`, and no test
of its own since `4756615e28` replaced
`register_output_stream_once_before_execution` with the shared
`register_before_execution`. `register_shuffle_partition_pusher` is still worth
keeping as a function because its test calls it. Could `setOutputStream` call
`register_before_execution` directly, and could the module doc at the top drop
`register_output_stream` from its list of core functions? The
once-and-before-execution rules stay covered by the pusher test, which goes
through the same helper.
##########
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:
Both callers build a complete `Operator` only for `writeFile` to take it
apart again. `writeNatively` and `CometNativeWriteExec` each run
`getParquetWriter.toBuilder ... build()` and then
`toBuilder.setParquetWriter(...).build()`, and `writeFile` repeats both steps
to set the path and the flag. Could `writeFile` take the template `Operator`
together with the caller's `ParquetWriter.Builder`, and set `output_path` and
`jvm_output_stream` on that builder? Each task would then build its plan once,
and "set up for the task except for the file" would show in the parameter types
rather than in the doc comment.
This helper also constructs `CometExecIterator` by hand, where both callers
used `CometExec.getCometIterator` before. Adding `outputStream:
Option[HadoopFileOutput] = None` to the existing overload would keep iterator
construction in `CometExec`.
##########
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:
Since #7090, `parquet.md` has a "Writing" section that holds the native
writer's scope (non-partitioned and non-bucketed writes only, and not with
`spark.sql.files.maxRecordsPerFile`), and the scan rows in this table link to
`parquet.md#reading`. Could this paragraph move into that section, next to the
scope sentence, with the `WriteFilesExec` and `DataWritingCommandExec` rows
linking to `parquet.md#writing`? Someone checking which destinations the writer
supports would then find them with the rest of its limits.
--
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]
