comphead commented on code in PR #6926:
URL: https://github.com/apache/datafusion-comet/pull/6926#discussion_r4543982361
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala:
##########
@@ -238,62 +239,44 @@ case class CometIcebergWriteExec(
child.getClass.getName)
}
- val childRDD = schedulerProbe
- .map(probe =>
- new IcebergSchedulerProbeRDD(
- inputRDD,
- probe,
- nativeOp.getIcebergWrite.getCommon.getDataLocation,
- nativeOp.getIcebergWrite.getCommon.getOperationId))
- .getOrElse(inputRDD)
- val numPartitions = childRDD.getNumPartitions
- // Native side emits the per-task Avro `DataFile` blob and the locations
it wrote, one Binary
- // column each (see `build_output_schema` in `iceberg_write.rs`).
- val numOutputCols = 2
- val capturedNativeOp = nativeOp
-
- childRDD.mapPartitionsInternal { iter =>
- val partitionId = TaskContext.getPartitionId()
- val taskAttemptId = TaskContext.get().taskAttemptId()
-
- // No child metric nodes: the write's native tree (IcebergWrite ->
Projection -> Scan) does
- // not line up with the JVM child operators, which own their own native
blocks; zipping
- // them by index would overwrite the child operators' SQLMetrics with
the write block's.
- // The root node still receives the native operator's own metrics by
name (`write_time`);
- // the file/row/byte SQL metrics and the task OutputMetrics are
populated in doExecute from
- // the decoded manifest instead.
- val nativeMetrics = CometMetricNode(metrics, Nil)
-
- val taskNativeOp = {
- val builder = capturedNativeOp.getIcebergWrite.toBuilder
- .setPartitionId(partitionId)
- .setTaskAttemptId(taskAttemptId)
- val icebergWrite = builder.build()
- capturedNativeOp.toBuilder.setIcebergWrite(icebergWrite).build()
+ // The native plan has one leaf, which reads `child` (see `buildInputScan`
in
+ // `CometIcebergNativeWrite`). A `ShuffleScan` leaf takes the child's
shuffle RDD and decodes
+ // its blocks natively; a `Scan` leaf imports the child's batches as an
Arrow stream.
+ val shuffleScanIndices = CometExec.findShuffleScanIndices(nativeOp)
+ val input =
+ if (shuffleScanIndices.nonEmpty) {
+ // `resolveInputObjects` needs the shuffle RDD itself, which the
scheduler probe's RDD
+ // would hide, so the probe's `beforeInput` runs in
`CometIcebergWriteRDD.compute`, ahead
+ // of the shuffle fetch that `resolveInputObjects` starts.
+ inputRDD
+ } else {
+ val probed = schedulerProbe
+ .map(probe =>
+ new IcebergSchedulerProbeRDD(
+ inputRDD,
+ probe,
+ nativeOp.getIcebergWrite.getCommon.getDataLocation,
+ nativeOp.getIcebergWrite.getCommon.getOperationId))
+ .getOrElse(inputRDD)
+ CometArrowStream.wrapColumnarBatchRDD(
+ probed,
+ CometUtils.fromAttributes(child.output),
+ CometArrowStream.NATIVE_TIMEZONE,
+ "CometIcebergWriteExec")
}
- val size = taskNativeOp.getSerializedSize
- val planBytes = new Array[Byte](size)
- val codedOutput = CodedOutputStream.newInstance(planBytes)
- taskNativeOp.writeTo(codedOutput)
- codedOutput.checkNoSpaceLeft()
-
- val execIterator = new CometExecIterator(
- CometExec.newIterId,
- CometArrowStream.inputObjects(
- iter,
- CometUtils.fromAttributes(child.output),
- "CometIcebergWriteExec"),
- numOutputCols,
- planBytes,
- nativeMetrics,
- numPartitions,
- partitionId,
- None,
- Seq.empty)
-
- execIterator
- }
+ // No child metric nodes: the write's native tree (IcebergWrite ->
Projection -> Scan or
+ // ShuffleScan) does not line up with the JVM child operators, which own
their own native
+ // blocks; zipping them by index would overwrite the child operators'
SQLMetrics with the
+ // write block's. The root node still receives the native operator's own
metrics by name
+ // (`write_time`); the file/row/byte SQL metrics and the task
OutputMetrics are populated in
+ // doExecute from the decoded manifest instead.
+ new CometIcebergWriteRDD(
+ input,
+ shuffleScanIndices,
+ nativeOp,
+ CometMetricNode(metrics, Nil),
+ if (shuffleScanIndices.nonEmpty) schedulerProbe else None)
Review Comment:
`shuffleScanIndices.nonEmpty` picks both the input RDD above and the probe
argument here. Could one branch return both, so the two choices cannot drift
apart? I have not compiled this:
```scala
val (input, shuffleProbe): (RDD[_], Option[IcebergSchedulerTestProbe]) =
if (shuffleScanIndices.nonEmpty) {
(inputRDD, schedulerProbe)
} else {
val probed = schedulerProbe.map(...).getOrElse(inputRDD)
(CometArrowStream.wrapColumnarBatchRDD(probed, ...), None)
}
```
--
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]
