andygrove commented on code in PR #6926:
URL: https://github.com/apache/datafusion-comet/pull/6926#discussion_r4529694134
##########
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:
You are right: with the wrapper kept, `resolveInputObjects` throws "expected
CometShuffledBatchRDD". I merged `main`, and a write with a `ShuffleScan` leaf
now takes the shuffle RDD as it is. `CometIcebergWriteRDD.compute` calls
`probe.beforeInput()` at the top, before `resolveInputObjects` starts the
fetch. A child that is read over FFI still goes through
`IcebergSchedulerProbeRDD`. The new test `the scheduler probe runs for a write
that reads a shuffle directly` repartitions to 4 partitions with AQE on and
coalescing off, checks the `SHUFFLE_SCAN` leaf, and asserts that the probe saw
all four partitions. It fails when I remove the call, and the other scheduler
probe tests still pass.
##########
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:
Added both. The `coalesced` case asserts that the final plan has an
`AQEShuffleReadExec` with a coalesced partition. A new `empty_partitions` case
does `REPARTITION(64, id)` with `WHERE id < 5` and turns coalescing off, so
most of the 64 `ShuffleScan`s reach the end of their input without a block
while the writer still emits its payload row. It checks the `SHUFFLE_SCAN`
leaf, the rows, and that the file record counts equal those of the FFI path.
##########
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:
Added date, timestamp, timestamp_ntz, decimal(12, 3), binary, boolean and
float columns. Nulls are in every one, and the double and float columns have
NaN, -0.0, both infinities and nulls. Every case is still written with direct
read on and off and compared with the source and with each other. All of it
passes at 5fa47b2116.
--
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]
