diff --git a/docs/source/contributor-guide/iceberg-writes.md b/docs/source/contributor-guide/iceberg-writes.md index 0fad3c35370..7ae9246e62e 100644 --- a/docs/source/contributor-guide/iceberg-writes.md +++ b/docs/source/contributor-guide/iceberg-writes.md @@ -125,6 +125,14 @@ batches from its child over FFI, so the conversion is declined unless the child native operator. This is why a write fed by a `LocalTableScanExec` (`INSERT ... VALUES`, a local DataFrame) stays on the JVM writer unless `spark.comet.exec.localTableScan.enabled` is also set. +A Comet shuffle directly under the write is the exception to the FFI path. When the child is a +sink that reads the shuffle directly (see +[Direct Read](native_shuffle.md#direct-read-shufflescan)), `buildInputScan` takes the sink's +`ShuffleScan` as the leaf of the write's native plan, so the shuffle blocks are decoded inside that +plan instead of in the JVM. Under AQE this covers a write over a shuffle stage or a coalesced read +of one. That is the plan for a repartitioned insert, and for a hash-distributed write that needs no +sort, such as one to a partitioned table without a sort order. + `CometIcebergWriteExec` is row-based (it emits the commit message row) over a columnar child, so Spark inserts a columnar-to-row transition beneath it. `EliminateRedundantTransitions` removes that transition in a dedicated pass that runs before its main `transformUp`, and the scaladoc on @@ -204,9 +212,10 @@ on the driver. Almost all of it is the per-write `IcebergWriteCommon`: including its `HadoopFileIO` fallback, rather than the wrapper's Hadoop configuration. Failure to resolve that delegate causes a plan-time fallback. -The per-task `partition_id` and `task_attempt_id` are stamped onto a copy of the proto inside the -task closure in `CometIcebergWriteExec.doExecuteColumnar`. The native side refuses to run without -them, since defaulting them would make every task write the same file names. +The per-task `partition_id` and `task_attempt_id` are stamped onto a copy of the proto in +`CometIcebergWriteRDD.compute`, the RDD that `CometIcebergWriteExec.doExecuteColumnar` returns. The +native side refuses to run without them, since defaulting them would make every task write the same +file names. `catalog_properties` can carry credentials. `CometIcebergWriteExec.stringArgs` deliberately prints only the data location and writer mode, because the default `argString` would put the protobuf's @@ -225,9 +234,10 @@ reserves for metadata columns, so those writes stay on iceberg-java. ### Native execution The planner builds `IcebergWriteExec` (`native/core/src/execution/operators/iceberg_write.rs`) over -the FFI scan of the child's batches. `run_write_task` builds iceberg-rust's partitioning writer once -per task, and `PartitionWriterBuilder` builds the rest of the stack once per partition, with that -partition's Parquet writer properties: +the scan of the child: the FFI scan of its batches, or the `ShuffleScan` of a shuffle it reads +directly. `run_write_task` builds iceberg-rust's partitioning writer once per task, and +`PartitionWriterBuilder` builds the rest of the stack once per partition, with that partition's +Parquet writer properties: ```text UnpartitionedWriter | FanoutPartitions | ClusteredWriter diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index db2eb0d7fb3..f6a94b5d2f6 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -260,6 +260,11 @@ of a two-phase aggregate, comes back from re-planning as the node already planne serialized as a `Scan` before the stage existed. `CometExecRule` refreshes such a node: once its input is a sink that emits a `ShuffleScan`, that `ShuffleScan` replaces the stale `Scan` leaf. +The native Iceberg writer builds its own native plan rather than joining its child's block, and it +follows the same decision: when its child is a sink that emits a `ShuffleScan`, +`CometIcebergNativeWrite` takes that `ShuffleScan` as the leaf of the write's plan. See +[Iceberg Writes](iceberg-writes.md). + The two are not alternatives on failure. If any output type fails `QueryPlanSerde.supportedDataType`, `convertToShuffleScan` records the fallback reason `Unsupported data type for shuffle direct read` and returns `None`. It does not retry as a regular `Scan`, and retrying would not help, because diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 403faea5317..5db85f30644 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -909,8 +909,9 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } /** - * Whether `op` is a writer whose native plan builds its own `Scan` over its child, so that the - * child runs as a separate native block instead of inside that plan. + * Whether `op` is a writer whose native plan reads its child through a leaf of its own, an FFI + * `Scan` or the `ShuffleScan` of a shuffle it reads directly, so that the child runs as a + * separate native block instead of inside that plan. */ private def scansChildAsSeparateBlock(op: SparkPlan): Boolean = op match { case _: CometNativeWriteExec | _: CometIcebergWriteExec | _: CometWriteFilesExec => true diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala index b385668291d..0f000f3d219 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala @@ -902,12 +902,12 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { op: IcebergWriteExec, builder: Operator.Builder, childOp: Operator*): Option[OperatorOuterClass.Operator] = { - val _ = (builder, childOp) // unused: we synthesise our own FFI scan child below + val _ = builder // unused: the write builds its own operator below try { for { icebergWrite <- buildIcebergWriteProto(op) - ffiScan <- buildFfiScan(op) - writeChild <- dropNonDataColumns(op, ffiScan) + input <- buildInputScan(op, childOp) + writeChild <- dropNonDataColumns(op, input) } yield OperatorOuterClass.Operator .newBuilder() .setPlanId(op.id) @@ -934,11 +934,11 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { * `decorate_batch_with_field_ids` rejects the batch. * * Plain writes already present only data columns, so no extra projection is needed. For 4.x - * ReplaceData we splice a `Projection` proto between our `IcebergWrite` op and the FFI `Scan`, - * selecting the upstream attributes whose names match the Iceberg schema's columns. The - * JVM-side child stays at the original wide output, so its `executeColumnar()` still emits the - * wide batches the FFI scan declares; the projection then strips them inside the native runtime - * before the writer sees the data. + * ReplaceData we splice a `Projection` proto between our `IcebergWrite` op and the input scan + * (`buildInputScan`), selecting the upstream attributes whose names match the Iceberg schema's + * columns. The JVM-side child stays at the original wide output, so its batches or shuffle + * blocks still carry the wide rows the scan declares; the projection then strips them inside + * the native runtime before the writer sees the data. */ private def dropNonDataColumns( op: IcebergWriteExec, @@ -996,6 +996,21 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { .build()) } + /** + * The leaf of the write's native plan, which reads its child. A child that reads a Comet + * shuffle directly is a sink whose native operator is a `ShuffleScan` (see + * `CometExchangeSink`). The write takes that `ShuffleScan` as its own leaf, so the shuffle + * blocks are decoded inside the write's plan instead of in the JVM, and are not exported and + * imported again over FFI. Any other child is read through an FFI `Scan` of its batches. + */ + private def buildInputScan( + op: IcebergWriteExec, + childOp: Seq[Operator]): Option[OperatorOuterClass.Operator] = + childOp match { + case Seq(child) if child.hasShuffleScan => Some(child) + case _ => buildFfiScan(op) + } + private def buildFfiScan(op: IcebergWriteExec): Option[OperatorOuterClass.Operator] = { val scan = NativeWriteUtils.buildFfiScan(op.child, op.id) if (scan.isEmpty) { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala index 98ba9babe77..5a324be6ee6 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala @@ -21,7 +21,7 @@ package org.apache.spark.sql.comet import java.util.concurrent.atomic.AtomicReference -import org.apache.spark.TaskContext +import org.apache.spark.{Partition, TaskContext} import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.Attribute @@ -56,7 +56,8 @@ import org.apache.comet.serde.OperatorOuterClass.Operator * @param originalPlan * The JVM Iceberg write operator restored when Comet reverts a transition-heavy stage. * @param child - * Comet native child (must be a [[CometNativeExec]] so columnar batches flow through FFI). + * Comet native child. The native plan imports its Arrow batches over FFI, or, when the child is + * a Comet shuffle that the plan reads with a `ShuffleScan`, decodes its shuffle blocks. * @param output * Commit-message output attributes copied from the JVM Iceberg write plan. * @param batchWrite @@ -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) } /** @@ -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()) + val (inputObjects, shuffleBlockIterators) = + CometExecRDD.resolveInputObjects(Seq(input), Array(split), shuffleScanIndices, context) + val partitionId = context.partitionId() + val icebergWrite = nativeOp.getIcebergWrite.toBuilder + .setPartitionId(partitionId) + .setTaskAttemptId(context.taskAttemptId()) + .build() + val planBytes = + CometExec.serializeNativePlan(nativeOp.toBuilder.setIcebergWrite(icebergWrite).build()) + // 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 + new CometExecIterator( + CometExec.newIterId, + inputObjects, + numOutputCols, + planBytes, + nativeMetrics, + numPartitions, + partitionId, + shuffleBlockIterators = shuffleBlockIterators) + } + + override def clearDependencies(): Unit = { + super.clearDependencies() + input = null + } +} diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 6f752267a63..2caa37de4c4 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -22,7 +22,7 @@ package org.apache.comet import java.io.File import java.nio.file.{FileAlreadyExistsException, Files} import java.sql.Timestamp -import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, TimeUnit} +import java.util.concurrent.{ConcurrentHashMap, ConcurrentLinkedQueue, CountDownLatch, TimeUnit} import java.util.concurrent.atomic.AtomicReference import scala.collection.mutable @@ -46,11 +46,11 @@ import org.apache.spark.sql.DataFrame import org.apache.spark.sql.Row import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.Attribute -import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometMergeRowsExec, IcebergCommitExec, IcebergSchedulerTestProbe, IcebergWriteExec} +import org.apache.spark.sql.comet.{CometExec, CometIcebergWriteExec, CometMergeRowsExec, IcebergCommitExec, IcebergSchedulerTestProbe, IcebergWriteExec} import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog import org.apache.spark.sql.connector.write.{BatchWrite, DataWriterFactory, PhysicalWriteInfo, Write, WriterCommitMessage} import org.apache.spark.sql.execution.{ColumnarToRowTransition, LeafExecNode, SparkPlan} -import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEShuffleReadExec} import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.streaming.Trigger @@ -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, " + + "arr ARRAY, m MAP, 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)))), + 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))), + IF(id % 23 = 0, NULL, DATE_ADD(DATE '1970-01-01', CAST(id * 7 AS INT))), + IF(id % 29 = 0, NULL, TIMESTAMP_MICROS(id * 1000003)), + IF(id % 31 = 0, NULL, CAST(TIMESTAMP_MICROS(id * 999983) AS TIMESTAMP_NTZ)), + IF(id % 37 = 0, NULL, CAST(CAST(id AS DECIMAL(12, 3)) / 7 AS DECIMAL(12, 3))), + IF(id % 41 = 0, NULL, CAST(CONCAT('b', CAST(id AS STRING)) AS BINARY)), + IF(id % 43 = 0, NULL, id % 2 = 0), + CASE id % 10 WHEN 0 THEN CAST('NaN' AS FLOAT) WHEN 1 THEN CAST(-0.0D AS FLOAT) + WHEN 2 THEN CAST('Infinity' AS FLOAT) WHEN 3 THEN CAST('-Infinity' AS FLOAT) + WHEN 4 THEN NULL ELSE CAST(id AS FLOAT) / 3 END + FROM range(0, 2000)""") + + // (case, hint, filter, AQE on, the write's leaf when direct read is on, extra confs). + // REBALANCE runs on Comet's JVM columnar shuffle, which writes the same block format as + // native shuffle. The `empty_partitions` case keeps 64 reducer partitions, of which the + // filter leaves at most five with rows, so most `ShuffleScan`s reach the end of their input + // without a block while the writer still emits its payload row. + val cases = Seq( + ("stage", "REPARTITION(8, id)", "", true, "SHUFFLE_SCAN", Seq.empty[(String, String)]), + ("coalesced", "REPARTITION(id)", "", true, "SHUFFLE_SCAN", Seq.empty[(String, String)]), + ("columnar", "REBALANCE", "", true, "SHUFFLE_SCAN", Seq.empty[(String, String)]), + ( + "empty_partitions", + "REPARTITION(64, id)", + "WHERE id < 5", + true, + "SHUFFLE_SCAN", + Seq(SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false")), + ("no_aqe", "REPARTITION(8, id)", "", false, "SCAN", Seq.empty[(String, String)])) + cases.foreach { case (name, hint, filter, adaptive, directReadLeaf, extraConfs) => + val expected = + spark.sql(s"SELECT * FROM $catalog.$ns.shuffle_src $filter").collect().toSeq + val fileRecordCounts = Seq(true, false).map { directRead => + val table = s"shuffle_${name}_$directRead" + spark.sql(s"CREATE TABLE $catalog.$ns.$table ($columns) USING iceberg") + // `withSQLConf` returns Unit on Spark 3.x, so it wraps only the statement. + val snapshot = withNativeEnabled { + captureWrite(table) { + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_SHUFFLE_DIRECT_READ_ENABLED.key -> directRead.toString) ++ + extraConfs): _*) { + spark.sql(s"INSERT INTO $catalog.$ns.$table " + + s"SELECT /*+ $hint */ * FROM $catalog.$ns.shuffle_src $filter") + } + } + } + assert( + snapshot.snapshotDelta == 1L, + s"$table: expected 1 commit, got ${snapshot.snapshotDelta}") + val leaves = snapshot.plans.flatMap { plan => + collectWithSubqueries(plan) { case write: CometIcebergWriteExec => + CometExec.nativeLeaves(write.nativeOp).map(_.getOpStructCase.name) + } + } + val expectedLeaf = if (directRead) directReadLeaf else "SCAN" + assert( + leaves == Seq(Seq(expectedLeaf)), + s"$table: expected the write to read its input with $expectedLeaf, got $leaves. " + + s"Plans:\n${snapshot.plans.mkString("\n--\n")}") + // The case is only the coalesced one if AQE coalesced the reducer partitions. + if (name == "coalesced") { + val coalesced = snapshot.plans.flatMap { plan => + collectWithSubqueries(plan) { case read: AQEShuffleReadExec => + read.hasCoalescedPartition + } + } + assert( + coalesced.contains(true), + s"$table: expected a coalesced shuffle read. Plans:\n${snapshot.plans.mkString("\n--\n")}") + } + checkAnswer(spark.table(s"$catalog.$ns.$table"), expected) + spark + .sql(s"SELECT record_count FROM $catalog.$ns.$table.files ORDER BY record_count") + .collect() + .map(_.getLong(0)) + .toSeq + } + assert( + fileRecordCounts.distinct.length == 1, + s"$name: data file record counts differ between direct read on and off: " + + fileRecordCounts) + } + } + } + + // The scheduler probe's RDD cannot wrap a shuffle that the write reads directly, because the + // direct read needs the shuffle RDD itself. The probe still has to run before each task reads + // its input, and the shuffle fetch is part of reading it. + test( + "native acceleration: the scheduler probe runs for a write that reads a shuffle directly") { + val table = "scheduler_shuffle_direct" + withSchedulerProbeSource(table) { source => + ShuffleReadProbe.reset() + val snapshot = IcebergSchedulerTestProbe.withProbe(ShuffleReadProbe) { + withNativeEnabled { + captureWrite(table) { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_DIRECT_READ_ENABLED.key -> "true") { + spark.read + .parquet(source.getAbsolutePath) + .repartition(4, org.apache.spark.sql.functions.col("id")) + .writeTo(s"$catalog.$ns.$table") + .append() + } + } + } + } + val leaves = snapshot.plans.flatMap { plan => + collectWithSubqueries(plan) { case write: CometIcebergWriteExec => + CometExec.nativeLeaves(write.nativeOp).map(_.getOpStructCase.name) + } + } + assert(leaves == Seq(Seq("SHUFFLE_SCAN")), leaves) + assert(ShuffleReadProbe.partitions == Set(0, 1, 2, 3)) + assert(snapshot.snapshotDelta == 1L) + assertRows(table, 0 until 5000) + } + } + for (adaptive <- Seq(false, true)) { test(s"transition-heavy fallback preserves Iceberg writes with AQE=$adaptive") { assumeNativeAcceleration() @@ -5340,3 +5485,17 @@ private object ConflictGate { def releaseWrite(): Unit = writeReleased.countDown() } + +/** Records the partitions whose input the probe saw, in the one JVM of a local test. */ +private[comet] object ShuffleReadProbe extends IcebergSchedulerTestProbe { + private val seen = ConcurrentHashMap.newKeySet[Integer]() + + def reset(): Unit = seen.clear() + def partitions: Set[Int] = seen.asScala.map(_.intValue()).toSet + + override def beforeInput(): Unit = seen.add(TaskContext.getPartitionId()) + override def afterInput(rows: Long, dataLocation: String, operationId: String): Unit = {} + override def afterHandoff(locations: Seq[String]): Unit = {} + override def accepted(partition: Int, locations: Seq[String]): Unit = {} + override def committed(): Unit = {} +}