andygrove commented on code in PR #6926:
URL: https://github.com/apache/datafusion-comet/pull/6926#discussion_r4543770579
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala:
##########
@@ -427,3 +410,60 @@ object CometIcebergWriteExec {
}
}
}
+
+/**
+ * Runs one native Iceberg write task per partition of `input`, the RDD that
feeds the write
+ * plan's one leaf. Like `mapPartitions`, it has `input`'s partitions and a
one-to-one dependency
+ * on it, but `compute` does not iterate `input`. It hands the partition to
+ * [[CometExecRDD.resolveInputObjects]], as every native block does, which
gives a `ShuffleScan`
+ * leaf the raw shuffle blocks and a `Scan` leaf an Arrow stream. Iterating a
shuffle input would
+ * fetch and decode its blocks in the JVM instead.
+ *
+ * Each task runs a copy of `nativeOp` stamped with the task's `partition_id`
and
+ * `task_attempt_id`, which the native writer needs to name its files.
+ */
+private class CometIcebergWriteRDD(
+ var input: RDD[_],
+ shuffleScanIndices: Set[Int],
+ nativeOp: Operator,
+ nativeMetrics: CometMetricNode,
+ shuffleProbe: Option[IcebergSchedulerTestProbe])
+ extends RDD[ColumnarBatch](input) {
+
+ // Read here on the driver, where the input's partitions are known.
+ private val numPartitions = input.getNumPartitions
+
+ override protected def getPartitions: Array[Partition] =
firstParent[Any].partitions
+
+ override def compute(split: Partition, context: TaskContext):
Iterator[ColumnarBatch] = {
+ // Before the shuffle fetch that `resolveInputObjects` starts, as the
probe's RDD does for a
+ // child that is not read natively.
+ shuffleProbe.foreach(_.beforeInput())
Review Comment:
Added in be6c67ca59. The comment in `compute` says the probe is test-only
and that this path fires only `beforeInput`, not `afterInput`, because the JVM
never sees the rows of a shuffle that the native plan decodes. `afterHandoff`
still fires from the write's common path.
##########
spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala:
##########
@@ -1133,6 +1133,151 @@ 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>, dt DATE, ts TIMESTAMP, ntz
TIMESTAMP_NTZ, " +
+ "dec DECIMAL(12, 3), bin BINARY, b BOOLEAN, f FLOAT"
+ 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))),
Review Comment:
Added in be6c67ca59. The `s` column now has CJK strings (`日本語-N`) and
strings with an accented letter and a 4-byte emoji (`é😀-N`), beside the ASCII
ones and the nulls.
##########
spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala:
##########
@@ -1133,6 +1133,151 @@ 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>, dt DATE, ts TIMESTAMP, ntz
TIMESTAMP_NTZ, " +
+ "dec DECIMAL(12, 3), bin BINARY, b BOOLEAN, f FLOAT"
+ 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))),
+ CASE id % 11 WHEN 0 THEN NULL WHEN 1 THEN CAST('NaN' AS DOUBLE) WHEN
2 THEN -0.0D
+ WHEN 3 THEN CAST('Infinity' AS DOUBLE) WHEN 4 THEN
CAST('-Infinity' AS DOUBLE)
+ ELSE id / 3.0 END,
+ IF(id % 13 = 0, NULL,
+ NAMED_STRUCT('a', CAST(id AS INT), 'b', IF(id % 2 = 0, NULL,
CAST(id AS STRING)))),
Review Comment:
Added in be6c67ca59. Two columns: `nested ARRAY<STRUCT<x INT, y STRING>>`,
with a null element and a struct of nulls, and `deep STRUCT<k: MAP<STRING,
ARRAY<INT>>, v: INT>`, with null maps, null array elements and empty arrays.
##########
spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala:
##########
@@ -1133,6 +1133,151 @@ 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>, " +
Review Comment:
Added in be6c67ca59. The date column now has 1900-01-01, 1969-12-31 and
9999-12-31, and the timestamp and timestamp_ntz columns have 1900-01-01, one
microsecond before the epoch (timestamp only) and 9999-12-31 23:59:59.999999. A
new `non_utc_session` case runs the write under `America/Los_Angeles`. Every
case still checks direct read on and off against the source and against each
other, and the new test and the scheduler probe tests pass.
--
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]
