comphead commented on code in PR #6926:
URL: https://github.com/apache/datafusion-comet/pull/6926#discussion_r4488371669
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala:
##########
@@ -236,54 +237,28 @@ case class CometIcebergWriteExec(
child.getClass.getName)
}
- 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 icebergWrite = capturedNativeOp.getIcebergWrite.toBuilder
- .setPartitionId(partitionId)
- .setTaskAttemptId(taskAttemptId)
- .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) {
+ childRDD
+ } else {
+ CometArrowStream.wrapColumnarBatchRDD(
+ childRDD,
+ 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))
Review Comment:
On `main`, #6601 wraps `child.executeColumnar()` in
`IcebergSchedulerProbeRDD` and passes that on as `childRDD`. If the rebase
keeps that wrapper as the input for a `ShuffleScan` leaf, `resolveInputObjects`
throws "expected CometShuffledBatchRDD", and the probe is meant to fire before
the shuffle fetch, which here starts inside `resolveInputObjects`. Could the
shuffle RDD go through unwrapped for a `ShuffleScan` leaf, with
`probe.beforeInput()` called at the top of `CometIcebergWriteRDD.compute`? The
existing probe tests use `coalesce(1)` or AQE off, so I expect they would not
catch it, but I have not run them.
##########
spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala:
##########
@@ -1151,6 +1151,79 @@ class CometIcebergWriteActionSuite
}
}
+ // A Comet shuffle directly under the native write is decoded inside the
write's native plan by
+ // a `ShuffleScan`, rather than decoded in the JVM and imported again over
FFI. Under AQE the
+ // write's child is the shuffle stage or a coalesced read of it. Without
AQE, Comet reads no
+ // shuffle directly, so the write keeps its FFI `Scan`. Each case is written
twice, with direct
+ // read on and off, so the two input paths are checked against each other as
well as the source.
+ test("native acceleration: a write over a shuffle decodes the shuffle blocks
natively") {
+ assumeNativeAcceleration()
+ withIcebergCatalog { _ =>
+ val columns = "id BIGINT, s STRING, d DOUBLE, st STRUCT<a: INT, b:
STRING>, " +
+ "arr ARRAY<INT>, m MAP<STRING, INT>"
Review Comment:
The PR moves decoding to `cast_and_stamp_schema`, but these columns cover
only string, double, struct, array and map. Could they add date, timestamp,
timestamp_ntz, decimal, binary, boolean and float, with NaN, -0.0, infinities
and nulls?
##########
spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala:
##########
@@ -1151,6 +1151,79 @@ class CometIcebergWriteActionSuite
}
}
+ // A Comet shuffle directly under the native write is decoded inside the
write's native plan by
+ // a `ShuffleScan`, rather than decoded in the JVM and imported again over
FFI. Under AQE the
+ // write's child is the shuffle stage or a coalesced read of it. Without
AQE, Comet reads no
+ // shuffle directly, so the write keeps its FFI `Scan`. Each case is written
twice, with direct
+ // read on and off, so the two input paths are checked against each other as
well as the source.
+ test("native acceleration: a write over a shuffle decodes the shuffle blocks
natively") {
+ assumeNativeAcceleration()
+ withIcebergCatalog { _ =>
+ val columns = "id BIGINT, s STRING, d DOUBLE, st STRUCT<a: INT, b:
STRING>, " +
+ "arr ARRAY<INT>, m MAP<STRING, INT>"
+ spark.sql(s"CREATE TABLE $catalog.$ns.shuffle_src ($columns) USING
iceberg")
+ spark.sql(s"""
+ INSERT INTO $catalog.$ns.shuffle_src
+ SELECT id,
+ IF(id % 7 = 0, NULL, CONCAT('s', CAST(id AS STRING))),
+ IF(id % 11 = 0, NULL, id / 3.0),
+ IF(id % 13 = 0, NULL,
+ NAMED_STRUCT('a', CAST(id AS INT), 'b', IF(id % 2 = 0, NULL,
CAST(id AS STRING)))),
+ IF(id % 17 = 0, NULL, ARRAY(CAST(id AS INT), NULL)),
+ IF(id % 19 = 0, NULL, MAP(CAST(id AS STRING), CAST(id AS INT)))
+ FROM range(0, 2000)""")
+ val expected = spark.table(s"$catalog.$ns.shuffle_src").collect().toSeq
+
+ // (case, hint, AQE on, the write's leaf when direct read is on).
REBALANCE runs on Comet's
+ // JVM columnar shuffle, which writes the same block format as native
shuffle.
+ val cases = Seq(
+ ("stage", "REPARTITION(8, id)", true, "SHUFFLE_SCAN"),
+ ("coalesced", "REPARTITION(id)", true, "SHUFFLE_SCAN"),
+ ("columnar", "REBALANCE", true, "SHUFFLE_SCAN"),
+ ("no_aqe", "REPARTITION(8, id)", false, "SCAN"))
Review Comment:
Nothing asserts that the `coalesced` case actually has a coalesced
`AQEShuffleReadExec`, and every reducer partition has rows. Could there be a
case with empty reducer partitions, for example `REPARTITION(64, id)` with
`WHERE id < 5`, so a `ShuffleScan` reaches EOF with no block while the writer
still emits its payload row?
--
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]
