parthchandra commented on code in PR #6926:
URL: https://github.com/apache/datafusion-comet/pull/6926#discussion_r4542405988
##########
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:
In `CometIcebergWriteRDD.compute`, the probe only fires `beforeInput()`. Can
we add a one-line comment here (or in the class scaladoc stating the shuffle
path intentionally skips the other probe callbacks because the probe is
test-only
##########
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:
The st struct is flat (a INT, b STRING). Maybe add a nested complex column
— array-of-struct or struct-containing-map — so nested-type handling over the
direct-read path is exercised, not just top-level struct/array/map.
##########
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:
The s column is ASCII only (CONCAT('s', id)). Can we add at least one
non-ASCII / multibyte string value (e.g. CJK characters) so string transport
through the shuffle blocks is covered.
##########
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:
Dates/timestamps stay in a mid range. Should we add boundary values
(pre-1970 date, far-future timestamp, and a non-UTC session timezone for the
ts/ntz columns) to the generated data?
--
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]
