From 992eb6c7098fa80717375eb901232adf77ba8812 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 4 Sep 2026 13:59:22 -0600 Subject: [PATCH 1/3] fix: restore columnar transitions under the native Iceberg write `CometIcebergWriteExec` was tagged `ColumnarToRowTransition` to stop Spark inserting a transition between it and its Comet-native child. That trait does more than suppress one insertion: Spark's `ensureOutputsRowBased` returns a `ColumnarToRowTransition` node untouched, so the entire subtree below the write skipped the transition-insertion pass. The Iceberg copy-on-write rewrite plan feeds a Spark-columnar `BatchScan` into row-based joins and filters, so the missing `ColumnarToRow` failed every DELETE / UPDATE / MERGE at runtime with `ColumnarBatch cannot be cast to InternalRow`. AQE hid it because each stage gets its own insertion pass when it materialises. Drop the trait so Spark walks the subtree normally, and strip the transition it now inserts below the write in `EliminateRedundantTransitions`, mirroring the existing `CometNativeWriteExec` arm. Closes #5689 --- .../rules/EliminateRedundantTransitions.scala | 27 ++++++- .../sql/comet/CometIcebergWriteExec.scala | 20 +++-- .../comet/CometIcebergWriteActionSuite.scala | 80 ++++++++++++++++++- 3 files changed, 117 insertions(+), 10 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala index 45fc26caee9..2cadede3272 100644 --- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala +++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala @@ -22,7 +22,7 @@ package org.apache.comet.rules import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.util.sideBySide -import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometIcebergWriteExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometSparkToColumnarExec} import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.comet.shims.{MapInBatchInfo, ShimCometMapInBatch} import org.apache.spark.sql.execution.{ColumnarToRowExec, RowToColumnarExec, SparkPlan} @@ -91,6 +91,18 @@ case class EliminateRedundantTransitions(session: SparkSession) // Write should be final operation in the plan case ColumnarToRowExec(nativeWrite: CometNativeWriteExec) => nativeWrite + // `CometIcebergWriteExec` is row-based (it emits the serialised Iceberg commit message) but + // consumes Arrow batches from its child over FFI, so Spark inserts a columnar-to-row + // transition *underneath* it. Strip it so `doExecuteColumnar` sees the columnar child + // directly; `CometIcebergNativeWrite.requiresNativeChildren` already guarantees that child + // was Comet-native when the write was converted. + // + // The write deliberately does not tag itself as a `ColumnarToRowTransition` to suppress the + // insertion: Spark leaves such a node untouched, so the whole subtree below the write is + // never visited and the transitions the rest of that subtree needs are never inserted + // (https://github.com/apache/datafusion-comet/issues/5689). + case w: CometIcebergWriteExec => + stripColumnarToRow(w.child).map(child => w.withNewChildren(Seq(child))).getOrElse(w) case c @ ColumnarToRowExec(child) if hasCometNativeChild(child) => val op = createColumnarToRowExec(child) if (c.logicalLink.isEmpty) { @@ -169,6 +181,19 @@ case class EliminateRedundantTransitions(session: SparkSession) } } + /** + * Unwraps a columnar-to-row transition, returning the columnar child underneath it, or `None` + * when `plan` is not such a transition. `transformUp` visits children first, so a plain + * `ColumnarToRowExec` over a Comet source has usually already been rewritten to one of the + * Comet variants by the time a parent arm looks at it; all three are handled. + */ + private def stripColumnarToRow(plan: SparkPlan): Option[SparkPlan] = plan match { + case CometNativeColumnarToRowExec(child) => Some(child) + case CometColumnarToRowExec(child) => Some(child) + case ColumnarToRowExec(child) => Some(child) + case _ => None + } + /** * If the given plan is a Comet ColumnarToRow transition, returns the columnar child the Python * UDF operator can consume directly. By the time this rule runs the earlier 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 2cafe2cf49a..25e0c5873c0 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 @@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.expressions.UnsafeProjection import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.comet.util.{Utils => CometUtils} import org.apache.spark.sql.connector.write.{BatchWrite, WriterCommitMessage} -import org.apache.spark.sql.execution.{ColumnarToRowTransition, SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.types.BinaryType import org.apache.spark.sql.vectorized.ColumnarBatch @@ -68,13 +68,7 @@ case class CometIcebergWriteExec( @transient table: AnyRef, partitionSpecId: Int) extends CometNativeExec - with UnaryExecNode - // We consume Arrow batches (via FFI) and emit row-shaped commit messages, so we are a - // columnar-to-row transition. Without this trait Spark's - // `ApplyColumnarRulesAndInsertTransitions` wedges a `CometNativeColumnarToRowExec` between - // us and the Comet-native child, which would then fail `child.executeColumnar()` in - // `doExecuteColumnar`. - with ColumnarToRowTransition { + with UnaryExecNode { override def originalPlan: SparkPlan = child @@ -86,6 +80,16 @@ case class CometIcebergWriteExec( // Native exec emits a single Binary column; the surrounding command framework expects rows, so // the outer commit exec calls executeCollect on us. supportsColumnar = false keeps Spark from // inserting a ColumnarToRow that would clash with our (Nil-output-like) row contract. + // + // We do consume Arrow batches from `child` over FFI, so Spark's + // `ApplyColumnarRulesAndInsertTransitions` inserts a columnar-to-row transition *below* us + // (our child is columnar-only, we are row-based). `EliminateRedundantTransitions` strips that + // transition back off so `doExecuteColumnar` sees the Comet-native child directly. + // + // Tagging this node as a `ColumnarToRowTransition` would also stop Spark inserting it, but at + // the cost of correctness: `ensureOutputsRowBased` returns such a node untouched, so nothing + // below the write is visited and the transitions the rest of that subtree needs are never + // inserted. See https://github.com/apache/datafusion-comet/issues/5689. override def supportsColumnar: Boolean = false override def executeCollect(): Array[InternalRow] = { diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 66ea7216786..c6e336bc9c3 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -32,8 +32,9 @@ import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.Row import org.apache.spark.sql.comet.{CometIcebergWriteExec, IcebergCommitExec, IcebergWriteExec} import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog -import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.{ColumnarToRowTransition, SparkPlan} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType, StructField, StructType} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus} @@ -696,6 +697,58 @@ class CometIcebergWriteActionSuite } } + // https://github.com/apache/datafusion-comet/issues/5689: the Iceberg CoW rewrite plan mixes a + // Spark-columnar `BatchScan` with row-based joins/filters underneath the write, so the write's + // subtree needs Spark to insert columnar-to-row transitions inside it. With AQE on those + // transitions are re-inserted when each stage materialises, which is why every other Iceberg + // suite in this file misses the problem; with AQE off the plan gets exactly one insertion pass + // and a missed transition fails the task with + // `ColumnarBatch cannot be cast to InternalRow`. + test("native acceleration: ReplaceData (CoW DELETE) with AQE disabled") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + createTable( + warehouseDir, + "native_cow_delete_no_aqe", + partitionSpec = "PARTITIONED BY (region)", + properties = Some("'write.delete.mode'='copy-on-write'")) + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + coalesceInsert( + "native_cow_delete_no_aqe", + Seq((1, "us-east", 10.0), (2, "us-west", 20.0), (3, "eu", 30.0), (4, "us-east", 40.0))) + } + // The IN-subquery is what puts a row-based join between the columnar CoW scan and the + // write, so the subtree genuinely needs a transition rather than being Comet-native + // end to end. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val snapshot = withNativeEnabled { + captureWrite("native_cow_delete_no_aqe") { + spark.sql("DELETE FROM cat.db.native_cow_delete_no_aqe WHERE id IN " + + "(SELECT col1 FROM VALUES (2) AS t(col1))") + } + } + assert( + snapshot.snapshotDelta == 1L, + s"expected 1 commit via native path, got ${snapshot.snapshotDelta}") + val nativeExecs = snapshot.plans.flatMap { p => + collectWithSubqueries(p) { case e: CometIcebergWriteExec => e } + } + assert( + nativeExecs.nonEmpty, + "expected >= 1 CometIcebergWriteExec in captured plans, got 0. Plans:\n" + + snapshot.plans.mkString("\n--\n")) + // The transition Spark inserts below the write must have been stripped again, otherwise + // `doExecuteColumnar` has no columnar child to pull Arrow batches from. + assert( + nativeExecs.forall(_.child.supportsColumnar), + "the native write must sit directly on its columnar child, got " + + nativeExecs.map(_.child.nodeName).mkString(", ")) + snapshot.plans.foreach(assertColumnarContract) + } + assertRows("native_cow_delete_no_aqe", Seq(1, 3, 4)) + } + } + test("native acceleration: ReplaceData (CoW UPDATE)") { assumeNativeAcceleration() withIcebergCatalog { warehouseDir => @@ -1688,6 +1741,31 @@ class CometIcebergWriteActionSuite assume(icebergAvailable, "Iceberg not available in classpath") } + /** + * Every row-consuming operator must receive row-based input. Spark guarantees that by inserting + * `ColumnarToRow` transitions in `ApplyColumnarRulesAndInsertTransitions`; an operator that + * Comet rewrote in a way that skips the insertion pass shows up here as a row-based node with a + * columnar-only child, and would fail at runtime with a `ColumnarBatch cannot be cast to + * InternalRow` `ClassCastException` rather than at planning time. + * + * `CometIcebergWriteExec` is the one legitimate exception: it is row-based on the outside but + * pulls Arrow batches from its Comet-native child over FFI (see the class docstring), so the + * transition below it is deliberately stripped again by `EliminateRedundantTransitions`. + */ + private def assertColumnarContract(plan: SparkPlan): Unit = { + val violations = collectWithSubqueries(plan) { + case p + if !p.supportsColumnar && !p.isInstanceOf[ColumnarToRowTransition] && + !p.isInstanceOf[CometIcebergWriteExec] && + p.children.exists(c => c.supportsColumnar && !c.supportsRowBased) => + p + } + assert( + violations.isEmpty, + "row-based operators consuming columnar-only children (missing ColumnarToRow): " + + s"${violations.map(_.nodeName).mkString(", ")}. Plan:\n$plan") + } + /** * Flip [[CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED]] for the duration of `action`. * From 585ffe8ee6d5f0d61a33d0695e4efb6285f21342 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 4 Sep 2026 14:54:19 -0600 Subject: [PATCH 2/3] docs: record why the plain ColumnarToRowExec form must stay in the strip Under AQE the write's child is `ColumnarToRowExec(AQEShuffleReadExec)`, and the `hasCometNativeChild` arm never rewrites it to a Comet variant because `QueryStageExec` is a `LeafExecNode`, so the `op.exists(...)` walk cannot see the Comet exchange inside the stage. Comment only. --- .../comet/rules/EliminateRedundantTransitions.scala | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala index 2cadede3272..c5bf7db0a35 100644 --- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala +++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala @@ -183,9 +183,15 @@ case class EliminateRedundantTransitions(session: SparkSession) /** * Unwraps a columnar-to-row transition, returning the columnar child underneath it, or `None` - * when `plan` is not such a transition. `transformUp` visits children first, so a plain - * `ColumnarToRowExec` over a Comet source has usually already been rewritten to one of the - * Comet variants by the time a parent arm looks at it; all three are handled. + * when `plan` is not such a transition. + * + * All three forms have to be handled. `transformUp` visits children first, so the + * `hasCometNativeChild` arm above has usually already rewritten a `ColumnarToRowExec` over a + * Comet source into one of the Comet variants by the time a parent arm looks at it. It has not + * when the source is an `AQEShuffleReadExec` over a `ShuffleQueryStageExec`: `QueryStageExec` + * is a `LeafExecNode`, so the `op.exists(...)` walk cannot see the Comet exchange inside it and + * the arm misses. That is the shape the Iceberg write gets under AQE, so the plain node reaches + * here. */ private def stripColumnarToRow(plan: SparkPlan): Option[SparkPlan] = plan match { case CometNativeColumnarToRowExec(child) => Some(child) From bc6429faec5339c84815311b9187ae6a6e28ad84 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 5 Sep 2026 09:36:36 -0600 Subject: [PATCH 3/3] fix: strip the Iceberg write's input transition before transition cancellation Review feedback: the write-boundary strip ran as an arm inside the rule's `transformUp`, so the generic `ColumnarToRowExec(CometSparkToColumnarExec)` cancellation reached the write's child first and destroyed the Arrow bridge the write's FFI input needs. Over a row source that arm drops the `CometSparkToColumnarExec` outright, leaving the write with a row child; over a Spark-columnar source it keeps a transition but hands the write Spark `ColumnarVector`s where the FFI adapter requires `CometVector`s. Both shapes are reachable through `spark.comet.sparkToColumnar.enabled`, whose `CometScanWrapper` passes `requiresNativeChildren` and is then unwrapped by `CometExecRule`, leaving the bridge directly beneath the write. Move the strip into its own pass ahead of `transformUp` so the write's input transition is removed before anything else can consume it. Tests: - Two plan-level tests in `CometIcebergWriteDetectionSuite` pin the boundary for both source representations. Both fail without this change. - `captureWrite` now runs `assertColumnarContract` on every captured plan rather than only the one test written to exercise it. - AQE-off CoW UPDATE and MERGE tests cover the two remaining rewrite shapes from #5689. The partitioned MERGE engages natively (the rebalance exchange makes the write's child Comet-native even though `MergeRowsExec` is not), so its subtree mixes a Spark-columnar scan, a row-based operator and Comet operators. --- .../rules/EliminateRedundantTransitions.scala | 60 ++++++++++----- .../comet/CometIcebergWriteActionSuite.scala | 72 +++++++++++++++++- .../CometIcebergWriteDetectionSuite.scala | 75 ++++++++++++++++++- 3 files changed, 185 insertions(+), 22 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala index c5bf7db0a35..6bb3f0dcd59 100644 --- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala +++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala @@ -74,7 +74,7 @@ case class EliminateRedundantTransitions(session: SparkSession) } private def _apply(plan: SparkPlan): SparkPlan = { - val eliminatedPlan = plan transformUp { + val eliminatedPlan = stripIcebergWriteInputTransition(plan) transformUp { case ColumnarToRowExec(shuffleExchangeExec: CometShuffleExchangeExec) if plan.conf.adaptiveExecutionEnabled => shuffleExchangeExec @@ -91,18 +91,6 @@ case class EliminateRedundantTransitions(session: SparkSession) // Write should be final operation in the plan case ColumnarToRowExec(nativeWrite: CometNativeWriteExec) => nativeWrite - // `CometIcebergWriteExec` is row-based (it emits the serialised Iceberg commit message) but - // consumes Arrow batches from its child over FFI, so Spark inserts a columnar-to-row - // transition *underneath* it. Strip it so `doExecuteColumnar` sees the columnar child - // directly; `CometIcebergNativeWrite.requiresNativeChildren` already guarantees that child - // was Comet-native when the write was converted. - // - // The write deliberately does not tag itself as a `ColumnarToRowTransition` to suppress the - // insertion: Spark leaves such a node untouched, so the whole subtree below the write is - // never visited and the transitions the rest of that subtree needs are never inserted - // (https://github.com/apache/datafusion-comet/issues/5689). - case w: CometIcebergWriteExec => - stripColumnarToRow(w.child).map(child => w.withNewChildren(Seq(child))).getOrElse(w) case c @ ColumnarToRowExec(child) if hasCometNativeChild(child) => val op = createColumnarToRowExec(child) if (c.logicalLink.isEmpty) { @@ -181,17 +169,49 @@ case class EliminateRedundantTransitions(session: SparkSession) } } + /** + * `CometIcebergWriteExec` is row-based (it emits the serialised Iceberg commit message) but + * consumes Arrow batches from its child over FFI, so Spark's + * `ApplyColumnarRulesAndInsertTransitions` inserts a columnar-to-row transition *underneath* + * it. Strip that transition so `doExecuteColumnar` sees the columnar child directly; + * `CometIcebergNativeWrite.requiresNativeChildren` already guarantees the child was + * Comet-native when the write was converted. + * + * The write deliberately does not tag itself as a `ColumnarToRowTransition` to suppress the + * insertion: Spark leaves such a node untouched, so the whole subtree below the write is never + * visited and the transitions the rest of that subtree needs are never inserted + * (https://github.com/apache/datafusion-comet/issues/5689). + * + * This runs as a separate pass *before* the main `transformUp`, not as an arm inside it, + * because the write's input transition has to be removed before the generic transition + * cancellation can consume it. `transformUp` visits children first, so the + * `ColumnarToRowExec(CometSparkToColumnarExec)` arm would otherwise rewrite the write's child + * before the write itself is visited, and that arm is destructive at this boundary: + * - over a row source it drops the `CometSparkToColumnarExec` as well, leaving the write with + * a row child and no Arrow producer at all; + * - over a Spark-columnar source it keeps a `ColumnarToRowExec` but drops the Arrow bridge, + * so the write would be handed Spark `ColumnarVector`s where the FFI adapter requires + * `CometVector`s. That shape is reachable whenever `spark.comet.sparkToColumnar.enabled` + * admits the write's source: `CometSparkToColumnarExec.createExec` wraps it in a + * `CometScanWrapper` (a `CometNativeExec`, so `requiresNativeChildren` accepts it), and + * `CometExecRule` then unwraps the placeholder, leaving the bridge directly beneath the + * write. + */ + private def stripIcebergWriteInputTransition(plan: SparkPlan): SparkPlan = plan.transform { + case w: CometIcebergWriteExec => + stripColumnarToRow(w.child).map(child => w.withNewChildren(Seq(child))).getOrElse(w) + } + /** * Unwraps a columnar-to-row transition, returning the columnar child underneath it, or `None` * when `plan` is not such a transition. * - * All three forms have to be handled. `transformUp` visits children first, so the - * `hasCometNativeChild` arm above has usually already rewritten a `ColumnarToRowExec` over a - * Comet source into one of the Comet variants by the time a parent arm looks at it. It has not - * when the source is an `AQEShuffleReadExec` over a `ShuffleQueryStageExec`: `QueryStageExec` - * is a `LeafExecNode`, so the `op.exists(...)` walk cannot see the Comet exchange inside it and - * the arm misses. That is the shape the Iceberg write gets under AQE, so the plain node reaches - * here. + * The plain `ColumnarToRowExec` is what Spark's insertion pass produces and is the only form + * observed in the suites. The two Comet variants are handled as well because this rule is part + * of `postColumnarTransitions` and AQE applies those rules once per materialised stage plus + * once for the final plan, so it can be handed a plan whose transitions an earlier pass already + * rewrote. Missing a variant would not be caught at planning time -- the write would fail at + * runtime with `requires a columnar (Comet native) child`. */ private def stripColumnarToRow(plan: SparkPlan): Option[SparkPlan] = plan match { case CometNativeColumnarToRowExec(child) => Some(child) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index c6e336bc9c3..8095e6feb05 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -743,7 +743,6 @@ class CometIcebergWriteActionSuite nativeExecs.forall(_.child.supportsColumnar), "the native write must sit directly on its columnar child, got " + nativeExecs.map(_.child.nodeName).mkString(", ")) - snapshot.plans.foreach(assertColumnarContract) } assertRows("native_cow_delete_no_aqe", Seq(1, 3, 4)) } @@ -805,6 +804,73 @@ class CometIcebergWriteActionSuite } } + // The remaining two CoW shapes from issue #5689, with AQE off so the plan gets exactly one + // transition-insertion pass. `captureWrite` runs `assertColumnarContract` on every captured + // plan, so these pin the same invariant as the DELETE case over the UPDATE and MERGE rewrite + // shapes, which put different operators between the columnar CoW scan and the write. + test("native acceleration: ReplaceData (CoW UPDATE) with AQE disabled") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + createTable( + warehouseDir, + "native_cow_update_no_aqe", + partitionSpec = "PARTITIONED BY (region)", + properties = Some("'write.update.mode'='copy-on-write'")) + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + coalesceInsert( + "native_cow_update_no_aqe", + Seq((1, "us-east", 10.0), (2, "us-west", 20.0), (3, "eu", 30.0))) + } + // As in the DELETE case, the IN-subquery is what puts a row-based join between the + // Spark-columnar CoW scan and the write. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + assertNativeWriteEngages("native_cow_update_no_aqe", Seq(1, 2, 3)) { + spark.sql( + "UPDATE cat.db.native_cow_update_no_aqe SET amount = amount * 2 " + + "WHERE id IN (SELECT col1 FROM VALUES (2) AS t(col1))") + } + } + val r = spark + .sql("SELECT amount FROM cat.db.native_cow_update_no_aqe WHERE id = 2") + .collect() + assert(r.length == 1 && r(0).getDouble(0) == 40.0, s"got ${r.toSeq}") + } + } + + test("native acceleration: ReplaceData (CoW MERGE) with AQE disabled") { + // Unlike the unpartitioned MERGE above, this one *does* engage natively: partitioning the + // table puts a `CometColumnarExchange` (REBALANCE_PARTITIONS_BY_COL) between the JVM + // `MergeRowsExec` and the write, so the write's own child is Comet-native and + // `requiresNativeChildren` is satisfied even though `MergeRowsExec` itself is not. + // + // That makes this the strongest of the three CoW shapes for issue #5689: the subtree below + // the write mixes a Spark-columnar `BatchScan`, a row-based `MergeRowsExec` and Comet + // operators, so it needs transitions inserted in three different places. + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + createTable( + warehouseDir, + "native_cow_merge_no_aqe", + partitionSpec = "PARTITIONED BY (region)", + properties = Some("'write.merge.mode'='copy-on-write'")) + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + coalesceInsert("native_cow_merge_no_aqe", Seq((1, "us-east", 10.0), (2, "us-west", 20.0))) + } + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + assertNativeWriteEngages("native_cow_merge_no_aqe", Seq(1, 2, 3)) { + spark.sql(""" + |MERGE INTO cat.db.native_cow_merge_no_aqe t + |USING (SELECT 2 AS id, 'us-west' AS region, 200.0 AS amount UNION ALL + | SELECT 3 AS id, 'eu' AS region, 30.0 AS amount) s + |ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET t.amount = s.amount + |WHEN NOT MATCHED THEN INSERT (id, region, amount) VALUES (s.id, s.region, s.amount) + |""".stripMargin) + } + } + } + } + test("native acceleration: complex types (struct, array, map) round-trip with field IDs") { assumeNativeAcceleration() withIcebergCatalog { _ => @@ -1686,6 +1752,10 @@ class CometIcebergWriteActionSuite private def captureWrite(tableName: String)(action: => Unit): WriteSnapshot = { val before = countSnapshots(tableName) val plans = capturePlans(spark)(action) + // Every write in this suite gets the columnar contract checked, not just the ones written to + // exercise it: a Comet rewrite that hides part of a subtree from Spark's transition-insertion + // pass is a whole class of bug (issue #5689) and the check is free once the plans are here. + plans.foreach(assertColumnarContract) WriteSnapshot(countSnapshots(tableName) - before, plans) } diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala index fc530a2d2e1..7c5631640c8 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala @@ -29,12 +29,20 @@ import org.apache.iceberg.hadoop.{HadoopConfigurable, HadoopFileIO} import org.apache.iceberg.io.{FileIO, InputFile, OutputFile} import org.apache.iceberg.util.SerializableSupplier import org.apache.spark.SparkConf +import org.apache.spark.rdd.RDD import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.comet.IcebergWriteExec +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} +import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometSparkToColumnarExec, IcebergWriteExec} +import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions, ColumnarToRowExec, LeafExecNode, SparkPlan} +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus import org.apache.comet.iceberg.IcebergReflection +import org.apache.comet.rules.EliminateRedundantTransitions import org.apache.comet.serde.{Compatible, SupportLevel, Unsupported} +import org.apache.comet.serde.OperatorOuterClass.Operator import org.apache.comet.serde.operator.CometIcebergNativeWrite class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTestBase { @@ -881,6 +889,71 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes fail(s"expected Unsupported for $tableName, got $other") } } + + /** + * Runs Spark's transition insertion followed by [[EliminateRedundantTransitions]] over a + * hand-built `CometIcebergWriteExec -> CometSparkToColumnarExec -> source` plan and returns the + * write's final child. + * + * Hand-built rather than driven through SQL because the shape depends on + * `spark.comet.sparkToColumnar.enabled` admitting the write's source operator, and the set of + * admitted operators is itself configurable. What matters is the rule's behaviour at that + * boundary, which this pins directly. + */ + private def writeChildAfterTransitionRules(source: SparkPlan): SparkPlan = { + val write = CometIcebergWriteExec( + Operator.newBuilder().build(), + CometSparkToColumnarExec(source), + batchWrite = null, + table = null, + partitionSpecId = 0) + val withTransitions = ApplyColumnarRulesAndInsertTransitions(Seq.empty, false).apply(write) + // Spark must insert a columnar-to-row transition below the row-based write; if it stops doing + // so the rest of the assertion is vacuous. + assert( + withTransitions.asInstanceOf[CometIcebergWriteExec].child.isInstanceOf[ColumnarToRowExec], + s"expected an inserted ColumnarToRowExec below the write, got:\n$withTransitions") + EliminateRedundantTransitions(spark) + .apply(withTransitions) + .asInstanceOf[CometIcebergWriteExec] + .child + } + + // https://github.com/apache/datafusion-comet/issues/5689: the write's input transition has to + // be stripped before the generic `ColumnarToRowExec(CometSparkToColumnarExec)` cancellation + // consumes it, otherwise that arm removes the Arrow bridge the write's FFI input depends on. + // Both source representations are covered because the cancellation treats them differently: + // over a row source it drops the bridge outright, over a Spark-columnar source it keeps a + // transition but leaves the write reading Spark `ColumnarVector`s instead of `CometVector`s. + test("row source keeps its Arrow bridge under the native Iceberg write") { + val source = TransitionProbeLeaf(columnar = false) + val child = writeChildAfterTransitionRules(source) + assert( + child == CometSparkToColumnarExec(source), + s"expected the write to sit directly on CometSparkToColumnarExec, got:\n$child") + } + + test("Spark-columnar source keeps its Arrow bridge under the native Iceberg write") { + val source = TransitionProbeLeaf(columnar = true) + val child = writeChildAfterTransitionRules(source) + assert( + child == CometSparkToColumnarExec(source), + s"expected the write to sit directly on CometSparkToColumnarExec, got:\n$child") + } +} + +/** + * Planning-only leaf used by the transition-boundary tests: `columnar` selects between the two + * source representations that can sit under a `CometSparkToColumnarExec`. Never executed. + */ +case class TransitionProbeLeaf(columnar: Boolean) extends LeafExecNode { + override def output: Seq[Attribute] = Seq( + AttributeReference("id", IntegerType, nullable = false)()) + override def supportsColumnar: Boolean = columnar + override protected def doExecute(): RDD[InternalRow] = + throw new UnsupportedOperationException("planning-only node") + override protected def doExecuteColumnar(): RDD[ColumnarBatch] = + throw new UnsupportedOperationException("planning-only node") } /**